package com.comet.opik.api.resources.v1.priv;

import com.comet.opik.api.BatchDelete;
import com.comet.opik.api.Column;
import com.comet.opik.api.Comment;
import com.comet.opik.api.Dataset;
import com.comet.opik.api.DatasetIdentifier;
import com.comet.opik.api.DatasetItem;
import com.comet.opik.api.DatasetItemBatch;
import com.comet.opik.api.DatasetItemBatchUpdate;
import com.comet.opik.api.DatasetItemSource;
import com.comet.opik.api.DatasetItemStreamRequest;
import com.comet.opik.api.DatasetItemUpdate;
import com.comet.opik.api.DatasetItemsDelete;
import com.comet.opik.api.DatasetLastExperimentCreated;
import com.comet.opik.api.DatasetLastOptimizationCreated;
import com.comet.opik.api.DatasetUpdate;
import com.comet.opik.api.Experiment;
import com.comet.opik.api.ExperimentItem;
import com.comet.opik.api.ExperimentItemsBatch;
import com.comet.opik.api.ExperimentType;
import com.comet.opik.api.FeedbackScoreBatchContainer;
import com.comet.opik.api.FeedbackScoreItem;
import com.comet.opik.api.Optimization;
import com.comet.opik.api.PageColumns;
import com.comet.opik.api.PercentageValues;
import com.comet.opik.api.Project;
import com.comet.opik.api.ProjectStats.AvgValueStat;
import com.comet.opik.api.ProjectStats.CountValueStat;
import com.comet.opik.api.ProjectStats.PercentageValueStat;
import com.comet.opik.api.ProjectStats.ProjectStatItem;
import com.comet.opik.api.Prompt;
import com.comet.opik.api.PromptVersion;
import com.comet.opik.api.ReactServiceErrorResponse;
import com.comet.opik.api.ScoreSource;
import com.comet.opik.api.Span;
import com.comet.opik.api.Trace;
import com.comet.opik.api.Visibility;
import com.comet.opik.api.VisibilityMode;
import com.comet.opik.api.error.ErrorMessage;
import com.comet.opik.api.filter.DatasetField;
import com.comet.opik.api.filter.DatasetFilter;
import com.comet.opik.api.filter.DatasetItemField;
import com.comet.opik.api.filter.DatasetItemFilter;
import com.comet.opik.api.filter.ExperimentsComparisonFilter;
import com.comet.opik.api.filter.ExperimentsComparisonValidKnownField;
import com.comet.opik.api.filter.FieldType;
import com.comet.opik.api.filter.Filter;
import com.comet.opik.api.filter.Operator;
import com.comet.opik.api.resources.utils.AuthTestUtils;
import com.comet.opik.api.resources.utils.ClickHouseContainerUtils;
import com.comet.opik.api.resources.utils.ClientSupportUtils;
import com.comet.opik.api.resources.utils.DurationUtils;
import com.comet.opik.api.resources.utils.MigrationUtils;
import com.comet.opik.api.resources.utils.MySQLContainerUtils;
import com.comet.opik.api.resources.utils.RedisContainerUtils;
import com.comet.opik.api.resources.utils.StatsUtils;
import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils;
import com.comet.opik.api.resources.utils.TestUtils;
import com.comet.opik.api.resources.utils.WireMockUtils;
import com.comet.opik.api.resources.utils.resources.DatasetResourceClient;
import com.comet.opik.api.resources.utils.resources.ExperimentResourceClient;
import com.comet.opik.api.resources.utils.resources.OptimizationResourceClient;
import com.comet.opik.api.resources.utils.resources.PromptResourceClient;
import com.comet.opik.api.resources.utils.resources.SpanResourceClient;
import com.comet.opik.api.resources.utils.resources.TraceResourceClient;
import com.comet.opik.api.resources.utils.traces.TraceAssertions;
import com.comet.opik.api.sorting.Direction;
import com.comet.opik.api.sorting.SortableFields;
import com.comet.opik.api.sorting.SortingField;
import com.comet.opik.domain.DatasetDAO;
import com.comet.opik.domain.FeedbackScoreMapper;
import com.comet.opik.domain.SpanType;
import com.comet.opik.domain.stats.StatsMapper;
import com.comet.opik.extensions.DropwizardAppExtensionProvider;
import com.comet.opik.extensions.RegisterApp;
import com.comet.opik.podam.PodamFactoryUtils;
import com.comet.opik.utils.JsonUtils;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.BigIntegerNode;
import com.fasterxml.jackson.databind.node.BooleanNode;
import com.fasterxml.jackson.databind.node.DoubleNode;
import com.fasterxml.jackson.databind.node.IntNode;
import com.fasterxml.jackson.databind.node.NullNode;
import com.fasterxml.jackson.databind.node.TextNode;
import com.fasterxml.uuid.Generators;
import com.fasterxml.uuid.impl.TimeBasedEpochGenerator;
import com.github.tomakehurst.wiremock.client.WireMock;
import com.google.common.collect.ImmutableMap;
import com.redis.testcontainers.RedisContainer;
import jakarta.ws.rs.client.Entity;
import jakarta.ws.rs.client.WebTarget;
import jakarta.ws.rs.core.GenericType;
import jakarta.ws.rs.core.HttpHeaders;
import jakarta.ws.rs.core.MediaType;
import jakarta.ws.rs.core.Response;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.ListUtils;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.hc.core5.http.HttpStatus;
import org.awaitility.Awaitility;
import org.glassfish.jersey.client.ChunkedInput;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Named;
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInstance;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;
import org.testcontainers.clickhouse.ClickHouseContainer;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.lifecycle.Startables;
import org.testcontainers.mysql.MySQLContainer;
import ru.vyarus.dropwizard.guice.test.ClientSupport;
import ru.vyarus.dropwizard.guice.test.jupiter.ext.TestDropwizardAppExtension;
import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate;
import uk.co.jemos.podam.api.PodamFactory;
import uk.co.jemos.podam.api.PodamUtils;

import java.math.BigDecimal;
import java.math.BigInteger;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.Spliterator;
import java.util.Spliterators;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import static com.comet.opik.api.Column.ColumnType;
import static com.comet.opik.api.DatasetItem.DatasetItemPage;
import static com.comet.opik.api.FeedbackScoreBatchContainer.FeedbackScoreBatch;
import static com.comet.opik.api.FeedbackScoreItem.FeedbackScoreBatchItem;
import static com.comet.opik.api.Visibility.PRIVATE;
import static com.comet.opik.api.Visibility.PUBLIC;
import static com.comet.opik.api.resources.utils.ClickHouseContainerUtils.DATABASE_NAME;
import static com.comet.opik.api.resources.utils.CommentAssertionUtils.IGNORED_FIELDS_COMMENTS;
import static com.comet.opik.api.resources.utils.FeedbackScoreAssertionUtils.assertFeedbackScoresIgnoredFieldsAndSetThemToNull;
import static com.comet.opik.api.resources.utils.TestHttpClientUtils.FAKE_API_KEY_MESSAGE;
import static com.comet.opik.api.resources.utils.TestHttpClientUtils.NO_API_KEY_RESPONSE;
import static com.comet.opik.api.resources.utils.TestHttpClientUtils.UNAUTHORIZED_RESPONSE;
import static com.comet.opik.api.resources.utils.TestUtils.getIdFromLocation;
import static com.comet.opik.api.resources.utils.TestUtils.toURLEncodedQueryParam;
import static com.comet.opik.api.resources.utils.WireMockUtils.WireMockRuntime;
import static com.comet.opik.api.resources.v1.priv.OptimizationsResourceTest.OPTIMIZATION_IGNORED_FIELDS;
import static com.comet.opik.infrastructure.auth.RequestContext.SESSION_COOKIE;
import static com.comet.opik.infrastructure.auth.RequestContext.WORKSPACE_HEADER;
import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.WRITE;
import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
import static com.github.tomakehurst.wiremock.client.WireMock.matching;
import static com.github.tomakehurst.wiremock.client.WireMock.matchingJsonPath;
import static com.github.tomakehurst.wiremock.client.WireMock.okJson;
import static com.github.tomakehurst.wiremock.client.WireMock.post;
import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
import static java.util.stream.Collectors.flatMapping;
import static java.util.stream.Collectors.groupingBy;
import static java.util.stream.Collectors.mapping;
import static java.util.stream.Collectors.toCollection;
import static java.util.stream.Collectors.toList;
import static java.util.stream.Collectors.toMap;
import static java.util.stream.Collectors.toSet;
import static java.util.stream.Collectors.toUnmodifiableSet;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.params.provider.Arguments.arguments;

@TestInstance(TestInstance.Lifecycle.PER_CLASS)
@DisplayName("Dataset Resource Test")
@ExtendWith(DropwizardAppExtensionProvider.class)
class DatasetsResourceTest {

    private static final String BASE_RESOURCE_URI = "%s/v1/private/datasets";
    private static final String EXPERIMENT_RESOURCE_URI = "%s/v1/private/experiments";
    private static final String DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH = "/items/experiments/items";

    private static final String URL_TEMPLATE_TRACES = "%s/v1/private/traces";

    public static final String[] IGNORED_FIELDS_LIST = {"feedbackScores", "createdAt", "lastUpdatedAt", "createdBy",
            "lastUpdatedBy", "comments"};
    public static final String[] IGNORED_FIELDS_DATA_ITEM = {"createdAt", "lastUpdatedAt", "experimentItems",
            "createdBy", "lastUpdatedBy", "datasetId", "tags"};
    public static final String[] DATASET_IGNORED_FIELDS = {"id", "createdAt", "lastUpdatedAt", "createdBy",
            "lastUpdatedBy", "experimentCount", "mostRecentExperimentAt", "lastCreatedExperimentAt",
            "datasetItemsCount", "lastCreatedOptimizationAt", "mostRecentOptimizationAt", "optimizationCount",
            "status", "latestVersion"};

    public static final String API_KEY = UUID.randomUUID().toString();
    private static final String USER = UUID.randomUUID().toString();
    private static final String WORKSPACE_ID = UUID.randomUUID().toString();
    private static final String TEST_WORKSPACE = UUID.randomUUID().toString();

    private static final TimeBasedEpochGenerator GENERATOR = Generators.timeBasedEpochGenerator();

    private final RedisContainer REDIS = RedisContainerUtils.newRedisContainer();
    private final MySQLContainer MYSQL = MySQLContainerUtils.newMySQLContainer();
    private final GenericContainer<?> ZOOKEEPER_CONTAINER = ClickHouseContainerUtils.newZookeeperContainer();
    private final ClickHouseContainer CLICKHOUSE = ClickHouseContainerUtils.newClickHouseContainer(ZOOKEEPER_CONTAINER);

    private final WireMockRuntime wireMock;

    @RegisterApp
    private final TestDropwizardAppExtension APP;

    {
        Startables.deepStart(REDIS, MYSQL, CLICKHOUSE, ZOOKEEPER_CONTAINER).join();

        wireMock = WireMockUtils.startWireMock();

        var databaseAnalyticsFactory = ClickHouseContainerUtils.newDatabaseAnalyticsFactory(
                CLICKHOUSE, DATABASE_NAME);

        MigrationUtils.runMysqlDbMigration(MYSQL);
        MigrationUtils.runClickhouseDbMigration(CLICKHOUSE);

        APP = TestDropwizardAppExtensionUtils.newTestDropwizardAppExtension(
                MYSQL.getJdbcUrl(),
                databaseAnalyticsFactory,
                wireMock.runtimeInfo(),
                REDIS.getRedisURI());
    }

    private final PodamFactory factory = PodamFactoryUtils.newPodamFactory();

    private String baseURI;
    private ClientSupport client;
    private PromptResourceClient promptResourceClient;
    private ExperimentResourceClient experimentResourceClient;
    private DatasetResourceClient datasetResourceClient;
    private TraceResourceClient traceResourceClient;
    private SpanResourceClient spanResourceClient;
    private TransactionTemplate mySqlTemplate;
    private OptimizationResourceClient optimizationResourceClient;

    @BeforeAll
    void setUpAll(ClientSupport client, TransactionTemplate mySqlTemplate) {

        this.baseURI = TestUtils.getBaseUrl(client);
        this.client = client;
        this.mySqlTemplate = mySqlTemplate;

        ClientSupportUtils.config(client);

        mockTargetWorkspace(API_KEY, TEST_WORKSPACE, WORKSPACE_ID);

        promptResourceClient = new PromptResourceClient(client, baseURI, factory);
        experimentResourceClient = new ExperimentResourceClient(client, baseURI, factory);
        datasetResourceClient = new DatasetResourceClient(client, baseURI);
        traceResourceClient = new TraceResourceClient(this.client, baseURI);
        spanResourceClient = new SpanResourceClient(this.client, baseURI);
        optimizationResourceClient = new OptimizationResourceClient(client, baseURI, factory);
    }

    @AfterAll
    void tearDownAll() {
        wireMock.server().stop();
    }

    private void mockTargetWorkspace(String apiKey, String workspaceName, String workspaceId) {
        AuthTestUtils.mockTargetWorkspace(wireMock.server(), apiKey, workspaceName, workspaceId, USER);
    }

    private void mockSessionCookieTargetWorkspace(String sessionToken, String workspaceName,
            String workspaceId) {
        AuthTestUtils.mockSessionCookieTargetWorkspace(wireMock.server(), sessionToken, workspaceName, workspaceId,
                USER);
    }

    private UUID createAndAssert(Dataset dataset) {
        return createAndAssert(dataset, API_KEY, TEST_WORKSPACE);
    }

    private UUID createAndAssert(Dataset dataset, String apiKey, String workspaceName) {
        try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                .request()
                .accept(MediaType.APPLICATION_JSON_TYPE)
                .header(HttpHeaders.AUTHORIZATION, apiKey)
                .header(WORKSPACE_HEADER, workspaceName)
                .post(Entity.json(dataset))) {

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201);
            assertThat(actualResponse.hasEntity()).isFalse();

            var id = getIdFromLocation(actualResponse.getLocation());

            assertThat(id).isNotNull();
            assertThat(id.version()).isEqualTo(7);

            return id;
        }
    }

    private Dataset getAndAssertEquals(UUID id, Dataset expected, String workspaceName, String apiKey) {
        var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                .path(id.toString())
                .request()
                .header(HttpHeaders.AUTHORIZATION, apiKey)
                .header(WORKSPACE_HEADER, workspaceName)
                .get();

        assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
        var actualEntity = actualResponse.readEntity(Dataset.class);

        assertThat(actualEntity.id()).isEqualTo(id);
        assertThat(actualEntity).usingRecursiveComparison()
                .ignoringFields(DATASET_IGNORED_FIELDS)
                .isEqualTo(expected);

        assertThat(actualEntity.lastUpdatedBy()).isEqualTo(USER);
        assertThat(actualEntity.createdBy()).isEqualTo(USER);
        assertThat(actualEntity.createdAt()).isInThePast();
        assertThat(actualEntity.lastUpdatedAt()).isInThePast();
        assertThat(actualEntity.experimentCount()).isNotNull();
        assertThat(actualEntity.datasetItemsCount()).isNotNull();

        return actualEntity;
    }

    @Nested
    @DisplayName("Api Key Authentication:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class ApiKey {

        private final String fakeApikey = UUID.randomUUID().toString();
        private final String okApikey = UUID.randomUUID().toString();

        Stream<Arguments> credentials() {
            return Stream.of(
                    arguments(okApikey, true, null),
                    arguments(fakeApikey, false, UNAUTHORIZED_RESPONSE),
                    arguments("", false, NO_API_KEY_RESPONSE));
        }

        Stream<Arguments> publicCredentials() {
            return Stream.of(
                    arguments(okApikey, PRIVATE),
                    arguments(fakeApikey, PUBLIC),
                    arguments("", PUBLIC));
        }

        Stream<Arguments> getDatasetPublicCredentials() {
            return Stream.of(
                    arguments(okApikey, PRIVATE, 200),
                    arguments(okApikey, PUBLIC, 200),
                    arguments("", PRIVATE, 404),
                    arguments("", PUBLIC, 200),
                    arguments(fakeApikey, PRIVATE, 404),
                    arguments(fakeApikey, PUBLIC, 200));
        }

        @BeforeEach
        void setUp() {

            wireMock.server().stubFor(
                    post(urlPathEqualTo("/opik/auth"))
                            .withHeader(HttpHeaders.AUTHORIZATION, equalTo(fakeApikey))
                            .withRequestBody(matchingJsonPath("$.workspaceName", matching(".+")))
                            .willReturn(WireMock.unauthorized().withHeader("Content-Type", "application/json")
                                    .withJsonBody(JsonUtils.readTree(
                                            new ReactServiceErrorResponse(FAKE_API_KEY_MESSAGE,
                                                    401)))));
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("create dataset: when api key is present, then return proper response")
        void createDataset__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, boolean shouldSucceed,
                io.dropwizard.jersey.errors.ErrorMessage errorMessage) {
            var project = factory.manufacturePojo(Project.class).toBuilder()
                    .id(null)
                    .build();

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .post(Entity.json(project))) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(errorMessage);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Get dataset by id: when api key is present, then return proper response")
        void getDatasetById__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, Visibility visibility,
                int expectedCode) {

            Dataset dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .visibility(visibility)
                    .build();

            var id = createAndAssert(dataset);

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);
            mockGetWorkspaceIdByName(TEST_WORKSPACE, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedCode);
                assertThat(actualResponse.hasEntity()).isTrue();
                if (expectedCode == 200) {
                    var actualEntity = actualResponse.readEntity(Dataset.class);
                    assertThat(actualEntity.id()).isEqualTo(id);
                } else {
                    assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Get dataset by name: when api key is present, then return proper response")
        void getDatasetByName__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, Visibility visibility,
                int expectedCode) {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .visibility(visibility)
                    .build();

            var id = createAndAssert(dataset);

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);
            mockGetWorkspaceIdByName(TEST_WORKSPACE, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("retrieve")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .post(Entity.json(new DatasetIdentifier(dataset.name())))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedCode);
                assertThat(actualResponse.hasEntity()).isTrue();
                if (expectedCode == 200) {
                    var actualEntity = actualResponse.readEntity(Dataset.class);
                    assertThat(actualEntity.id()).isEqualTo(id);
                } else {
                    assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
                }
            }
        }

        @ParameterizedTest
        @MethodSource("publicCredentials")
        @DisplayName("Get datasets: when api key is present, then return proper response")
        void getDatasets__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, Visibility visibility) {

            String workspaceName = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(okApikey, workspaceName, workspaceId);
            mockGetWorkspaceIdByName(workspaceName, workspaceId);

            List<Dataset> expected = prepareDatasetsListWithOnePublic();

            expected.forEach(dataset -> createAndAssert(dataset, okApikey, workspaceName));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();
                var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);
                if (visibility == PRIVATE) {
                    assertThat(actualEntity.content()).hasSize(expected.size());
                } else {
                    assertThat(actualEntity.content()).hasSize(1);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("Update dataset: when api key is present, then return proper response")
        void updateDataset__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, boolean shouldSucceed,
                io.dropwizard.jersey.errors.ErrorMessage errorMessage) {

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            var id = createAndAssert(dataset);

            var update = factory.manufacturePojo(DatasetUpdate.class);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.json(update))) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(errorMessage);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("Delete dataset: when api key is present, then return proper response")
        void deleteDataset__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, boolean shouldSucceed,
                io.dropwizard.jersey.errors.ErrorMessage errorMessage) {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            var id = createAndAssert(dataset);

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .delete()) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(errorMessage);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("Create dataset items: when api key is present, then return proper response")
        void createDatasetItems__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, boolean shouldSucceed,
                io.dropwizard.jersey.errors.ErrorMessage errorMessage) {

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .map(item -> item.toBuilder()
                            .id(null)
                            .build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .put(Entity.json(batch))) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(errorMessage);
                }
            }

        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Get dataset items by dataset id: when api key is present, then return proper response")
        void getDatasetItemsByDatasetId__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey,
                Visibility visibility,
                int expectedCode) {

            var datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .visibility(visibility)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .map(item -> item.toBuilder()
                            .id(null)
                            .build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);
            mockGetWorkspaceIdByName(TEST_WORKSPACE, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path("items")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedCode);
                assertThat(actualResponse.hasEntity()).isTrue();
                if (expectedCode == 200) {
                    var actualEntity = actualResponse.readEntity(DatasetItemPage.class);
                    assertThat(actualEntity.content().size()).isEqualTo(items.size());
                } else {
                    assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Stream dataset items: when api key is present, then return proper response")
        void streamDatasetItems__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, Visibility visibility,
                int expectedCode) {
            String name = UUID.randomUUID().toString();

            var datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .name(name)
                    .visibility(visibility)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .map(item -> item.toBuilder()
                            .id(null)
                            .build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);
            mockGetWorkspaceIdByName(TEST_WORKSPACE, WORKSPACE_ID);

            var request = new DatasetItemStreamRequest(name, null, null);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("stream")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .accept(MediaType.APPLICATION_OCTET_STREAM)
                    .post(Entity.json(request))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                if (expectedCode == 200) {
                    assertThat(actualResponse.hasEntity()).isTrue();
                    List<DatasetItem> actualItems = getStreamedItems(actualResponse);
                    assertThat(actualItems.size()).isEqualTo(items.size());

                } else {
                    assertThat(actualResponse.hasEntity()).isFalse();
                }
            }
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("Delete dataset items: when api key is present, then return proper response")
        void deleteDatasetItems__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, boolean shouldSucceed,
                io.dropwizard.jersey.errors.ErrorMessage errorMessage) {

            var datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);

            var delete = new DatasetItemsDelete(items.stream().map(DatasetItem::id).toList());

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("delete")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .post(Entity.json(delete))) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(errorMessage);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Get dataset item by id: when api key is present, then return proper response")
        void getDatasetItemById__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey,
                Visibility visibility,
                int expectedCode) {
            String name = UUID.randomUUID().toString();

            var datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .name(name)
                    .visibility(visibility)
                    .build());

            var item = factory.manufacturePojo(DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(datasetId)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            mockTargetWorkspace(okApikey, TEST_WORKSPACE, WORKSPACE_ID);
            mockGetWorkspaceIdByName(TEST_WORKSPACE, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path(item.id().toString())
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedCode);
                assertThat(actualResponse.hasEntity()).isTrue();
                if (expectedCode == 200) {
                    var actualEntity = actualResponse.readEntity(DatasetItem.class);
                    assertThat(actualEntity.id()).isEqualTo(item.id());
                } else {
                    assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
                }
            }

        }

    }

    @Nested
    @DisplayName("Session Token Cookie Authentication:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class SessionTokenCookie {

        private final String sessionToken = UUID.randomUUID().toString();
        private final String fakeSessionToken = UUID.randomUUID().toString();

        @BeforeAll
        void setUp() {
            wireMock.server().stubFor(
                    post(urlPathEqualTo("/opik/auth-session"))
                            .withCookie(SESSION_COOKIE, equalTo(sessionToken))
                            .withRequestBody(matchingJsonPath("$.workspaceName", matching(".+")))
                            .willReturn(okJson(AuthTestUtils.newWorkspaceAuthResponse(USER, WORKSPACE_ID))));

            wireMock.server().stubFor(
                    post(urlPathEqualTo("/opik/auth-session"))
                            .withCookie(SESSION_COOKIE, equalTo(fakeSessionToken))
                            .withRequestBody(matchingJsonPath("$.workspaceName", matching(".+")))
                            .willReturn(WireMock.unauthorized().withHeader("Content-Type", "application/json")
                                    .withJsonBody(JsonUtils.readTree(
                                            new ReactServiceErrorResponse(FAKE_API_KEY_MESSAGE,
                                                    401)))));
        }

        Stream<Arguments> credentials() {
            return Stream.of(
                    arguments(sessionToken, true, "OK_" + UUID.randomUUID()),
                    arguments(fakeSessionToken, false, UUID.randomUUID().toString()));
        }

        Stream<Arguments> publicCredentials() {
            return Stream.of(
                    arguments(sessionToken, PRIVATE, "OK_" + UUID.randomUUID()),
                    arguments(fakeSessionToken, PUBLIC, UUID.randomUUID().toString()));
        }

        Stream<Arguments> getDatasetPublicCredentials() {
            return Stream.of(
                    arguments(sessionToken, PRIVATE, "OK_" + UUID.randomUUID(), 200),
                    arguments(sessionToken, PUBLIC, "OK_" + UUID.randomUUID(), 200),
                    arguments(fakeSessionToken, PRIVATE, UUID.randomUUID().toString(), 404),
                    arguments(fakeSessionToken, PUBLIC, UUID.randomUUID().toString(), 200));
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("create dataset: when session token is present, then return proper response")
        void createDataset__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                boolean shouldSucceed, String workspaceName) {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            mockSessionCookieTargetWorkspace(this.sessionToken, workspaceName, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .post(Entity.json(dataset))) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(201);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(UNAUTHORIZED_RESPONSE);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Get dataset by id: when session token is present, then return proper response")
        void getDatasetById__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                Visibility visibility,
                String workspaceName, int expectedCode) {

            Dataset dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .visibility(visibility)
                    .build();

            mockGetWorkspaceIdByName(workspaceName, WORKSPACE_ID);

            var id = createAndAssert(dataset);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedCode);
                assertThat(actualResponse.hasEntity()).isTrue();
                if (expectedCode == 200) {
                    var actualEntity = actualResponse.readEntity(Dataset.class);
                    assertThat(actualEntity.id()).isEqualTo(id);
                } else {
                    assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Get dataset by name: when session token is present, then return proper response")
        void getDatasetByName__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                Visibility visibility,
                String workspaceName, int expectedCode) {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .visibility(visibility)
                    .build();

            var id = createAndAssert(dataset);

            mockSessionCookieTargetWorkspace(this.sessionToken, workspaceName, WORKSPACE_ID);
            mockGetWorkspaceIdByName(workspaceName, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("retrieve")
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .post(Entity.json(new DatasetIdentifier(dataset.name())))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedCode);
                assertThat(actualResponse.hasEntity()).isTrue();

                if (expectedCode == 200) {
                    var actualEntity = actualResponse.readEntity(Dataset.class);
                    assertThat(actualEntity.id()).isEqualTo(id);
                } else {
                    assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
                }
            }
        }

        @ParameterizedTest
        @MethodSource("publicCredentials")
        @DisplayName("Get datasets: when session token is present, then return proper response")
        void getDatasets__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                Visibility visibility, String workspaceName) {

            String workspaceId = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);
            mockGetWorkspaceIdByName(workspaceName, workspaceId);
            mockSessionCookieTargetWorkspace(this.sessionToken, workspaceName, workspaceId);

            List<Dataset> expected = prepareDatasetsListWithOnePublic();

            expected.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();
                var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);
                if (visibility == PRIVATE) {
                    assertThat(actualEntity.content()).hasSize(expected.size());
                } else {
                    assertThat(actualEntity.content()).hasSize(1);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("Update dataset: when session token is present, then return proper response")
        void updateDataset__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                boolean shouldSucceed, String workspaceName) {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            var id = createAndAssert(dataset);

            var update = factory.manufacturePojo(DatasetUpdate.class);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .put(Entity.json(update))) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(UNAUTHORIZED_RESPONSE);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("Delete dataset: when session token is present, then return proper response")
        void deleteDataset__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                boolean shouldSucceed, String workspaceName) {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            var id = createAndAssert(dataset);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .delete()) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(UNAUTHORIZED_RESPONSE);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("Create dataset items: when session token is present, then return proper response")
        void createDatasetItems__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                boolean shouldSucceed, String workspaceName) {

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .map(item -> item.toBuilder()
                            .id(null)
                            .build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            mockSessionCookieTargetWorkspace(this.sessionToken, workspaceName, WORKSPACE_ID);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .put(Entity.json(batch))) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(UNAUTHORIZED_RESPONSE);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Get dataset items by dataset id: when session token is present, then return proper response")
        void getDatasetItemsByDatasetId__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                Visibility visibility,
                String workspaceName, int expectedCode) {
            mockGetWorkspaceIdByName(workspaceName, WORKSPACE_ID);

            var datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .visibility(visibility)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .map(item -> item.toBuilder()
                            .id(null)
                            .build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path("items")
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedCode);
                assertThat(actualResponse.hasEntity()).isTrue();

                if (expectedCode == 200) {
                    var actualEntity = actualResponse.readEntity(DatasetItemPage.class);
                    assertThat(actualEntity.content().size()).isEqualTo(items.size());
                } else {
                    assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Stream dataset items: when session token is present, then return proper response")
        void getDatasetItemsStream__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                Visibility visibility,
                String workspaceName, int expectedCode) {
            String name = UUID.randomUUID().toString();

            var datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .name(name)
                    .visibility(visibility)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .map(item -> item.toBuilder()
                            .id(null)
                            .build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            mockSessionCookieTargetWorkspace(this.sessionToken, workspaceName, WORKSPACE_ID);
            mockGetWorkspaceIdByName(workspaceName, WORKSPACE_ID);

            var request = new DatasetItemStreamRequest(name, null, null);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("stream")
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_OCTET_STREAM)
                    .post(Entity.json(request))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);

                if (expectedCode == 200) {
                    assertThat(actualResponse.hasEntity()).isTrue();
                    List<DatasetItem> actualItems = getStreamedItems(actualResponse);
                    assertThat(actualItems.size()).isEqualTo(items.size());

                } else {
                    assertThat(actualResponse.hasEntity()).isFalse();
                }
            }
        }

        @ParameterizedTest
        @MethodSource("credentials")
        @DisplayName("Delete dataset items: when session token is present, then return proper response")
        void deleteDatasetItems__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                boolean shouldSucceed, String workspaceName) {

            var datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            mockSessionCookieTargetWorkspace(this.sessionToken, workspaceName, WORKSPACE_ID);

            var delete = new DatasetItemsDelete(items.stream().map(DatasetItem::id).toList());

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("delete")
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .post(Entity.json(delete))) {

                if (shouldSucceed) {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                    assertThat(actualResponse.hasEntity()).isFalse();
                } else {
                    assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                    assertThat(actualResponse.hasEntity()).isTrue();
                    assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
                            .isEqualTo(UNAUTHORIZED_RESPONSE);
                }
            }
        }

        @ParameterizedTest
        @MethodSource("getDatasetPublicCredentials")
        @DisplayName("Get dataset item by id: when session token is present, then return proper response")
        void getDatasetItemById__whenSessionTokenIsPresent__thenReturnProperResponse(String sessionToken,
                Visibility visibility,
                String workspaceName, int expectedCode) {

            mockGetWorkspaceIdByName(workspaceName, WORKSPACE_ID);

            String name = UUID.randomUUID().toString();

            var datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .name(name)
                    .visibility(visibility)
                    .build());

            var item = factory.manufacturePojo(DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(datasetId)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path(item.id().toString())
                    .request()
                    .cookie(SESSION_COOKIE, sessionToken)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedCode);
                assertThat(actualResponse.hasEntity()).isTrue();
                if (expectedCode == 200) {
                    var actualEntity = actualResponse.readEntity(DatasetItem.class);
                    assertThat(actualEntity.id()).isEqualTo(item.id());
                } else {
                    assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
                }
            }

        }

    }

    @Nested
    @DisplayName("Create:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class CreateDataset {

        @Test
        @DisplayName("Success")
        void create__success() {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            createAndAssert(dataset);
        }

        @Test
        @DisplayName("when description is multiline, then accept the request")
        void create__whenDescriptionIsMultiline__thenAcceptTheRequest() {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .description("""
                            Test
                            Description
                            """)
                    .build();

            createAndAssert(dataset);
        }

        @Test
        @DisplayName("when creating datasets with same name in different workspaces, then accept the request")
        void create__whenCreatingDatasetsWithSameNameInDifferentWorkspaces__thenAcceptTheRequest() {

            var name = UUID.randomUUID().toString();
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset1 = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .name(name)
                    .build();

            var dataset2 = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .name(name)
                    .build();

            createAndAssert(dataset1, apiKey, workspaceName);
            createAndAssert(dataset2);
        }

        @Test
        @DisplayName("when description is null, then accept the request")
        void create__whenDescriptionIsNull__thenAcceptNameCreate() {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .description(null)
                    .build();

            createAndAssert(dataset);
        }

        private Stream<Arguments> invalidDataset() {
            return Stream.of(
                    arguments(factory.manufacturePojo(Dataset.class).toBuilder().name(null).build(),
                            "name must not be blank"),
                    arguments(factory.manufacturePojo(Dataset.class).toBuilder().name("").build(),
                            "name must not be blank"),
                    arguments(factory.manufacturePojo(Dataset.class).toBuilder().description("a".repeat(256)).build(),
                            "description cannot exceed 255 characters"));
        }

        @ParameterizedTest
        @MethodSource("invalidDataset")
        @DisplayName("when request is not valid, then return 422")
        void create__whenRequestIsNotValid__thenReturn422(Dataset dataset, String errorMessage) {

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI)).request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(dataset))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(errorMessage);
            }

        }

        @Test
        @DisplayName("when dataset name already exists, then reject the request")
        void create__whenDatasetNameAlreadyExists__thenRejectNameCreate() {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            createAndAssert(dataset);

            createAndAssertConflict(dataset, "Dataset already exists with name '%s'".formatted(dataset.name()));
        }

        @Test
        @DisplayName("when dataset id already exists, then reject the request")
        void create__whenDatasetIdAlreadyExists__thenRejectNameCreate() {

            var dataset = factory.manufacturePojo(Dataset.class);
            var dataset2 = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(dataset.id())
                    .build();

            createAndAssert(dataset);

            createAndAssertConflict(dataset2, "Dataset already exists with name '%s'".formatted(dataset2.name()));
        }

        private void createAndAssertConflict(Dataset dataset, String conflictMessage) {
            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI)).request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.entity(dataset, MediaType.APPLICATION_JSON_TYPE))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(409);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(conflictMessage);
            }
        }
    }

    @Nested
    @DisplayName("Get: {id, name}")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class GetDataset {

        @Test
        @DisplayName("Success")
        void getDatasetById() {
            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            var id = createAndAssert(dataset);

            getAndAssertEquals(id, dataset, TEST_WORKSPACE, API_KEY);
        }

        @Test
        @DisplayName("when dataset not found, then return 404")
        void getDatasetById__whenDatasetNotFound__whenReturn404() {

            var id = UUID.randomUUID().toString();

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get();

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
            assertThat(actualResponse.hasEntity()).isTrue();
            assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
        }

        @Test
        @DisplayName("when retrieving dataset by name, then return dataset")
        void getDatasetByIdentifier() {
            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            createAndAssert(dataset);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("retrieve")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(new DatasetIdentifier(dataset.name())))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualEntity = actualResponse.readEntity(Dataset.class);
                assertThat(actualEntity).usingRecursiveComparison()
                        .ignoringFields(DATASET_IGNORED_FIELDS)
                        .isEqualTo(dataset);
            }
        }

        @Test
        @DisplayName("when dataset not found by dataset name, then return 404")
        void getDatasetByIdentifier__whenDatasetItemNotFound__thenReturn404() {
            var name = UUID.randomUUID().toString();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("retrieve")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(new DatasetIdentifier(name)))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
            }
        }

        @Test
        @DisplayName("when dataset has optimizations linked to it, then return dataset with optimizations summary")
        void getDatasetById__whenDatasetHasOptimizationsLinkedToIt__thenReturnDatasetWithOptimizationSummary() {
            var dataset = factory.manufacturePojo(Dataset.class);
            createAndAssert(dataset);

            Instant beforeCreateOptimizations = Instant.now();
            int optimizationsCnt = 5;
            for (int i = 0; i < optimizationsCnt; i++) {
                var optimization = optimizationResourceClient.createPartialOptimization().datasetName(dataset.name())
                        .build();
                optimizationResourceClient.create(optimization, API_KEY, TEST_WORKSPACE);
            }

            var actualDataset = getAndAssertEquals(dataset.id(), dataset, TEST_WORKSPACE, API_KEY);

            assertThat(actualDataset.optimizationCount()).isEqualTo(optimizationsCnt);
            assertThat(actualDataset.mostRecentOptimizationAt()).isAfter(beforeCreateOptimizations);
            assertThat(actualDataset.lastCreatedOptimizationAt()).isAfter(beforeCreateOptimizations);
        }

        @Test
        @DisplayName("when dataset has experiments linked to it, then return dataset with experiment summary")
        void getDatasetById__whenDatasetHasExperimentsLinkedToIt__thenReturnDatasetWithExperimentSummary() {

            var dataset = factory.manufacturePojo(Dataset.class);

            createAndAssert(dataset);

            var experiment1 = experimentResourceClient.createPartialExperiment()
                    .datasetName(dataset.name())
                    .build();

            var experiment2 = experimentResourceClient.createPartialExperiment()
                    .datasetName(dataset.name())
                    .build();

            createAndAssert(experiment1, API_KEY, TEST_WORKSPACE);
            createAndAssert(experiment2, API_KEY, TEST_WORKSPACE);

            var datasetItems = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class);

            DatasetItemBatch batch = new DatasetItemBatch(dataset.name(), null, datasetItems);

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Creating two traces with input, output and scores
            var trace1 = factory.manufacturePojo(Trace.class);
            createTrace(trace1, API_KEY, TEST_WORKSPACE);

            var trace2 = factory.manufacturePojo(Trace.class);
            createTrace(trace2, API_KEY, TEST_WORKSPACE);

            var traces = List.of(trace1, trace2);

            // Creating 5 scores peach each of the two traces above
            var scores1 = PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItem.class)
                    .stream()
                    .map(feedbackScoreBatchItem -> feedbackScoreBatchItem.toBuilder()
                            .id(trace1.id())
                            .projectName(trace1.projectName())
                            .build())
                    .toList();

            var scores2 = PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItem.class)
                    .stream()
                    .map(feedbackScoreBatchItem -> feedbackScoreBatchItem.toBuilder()
                            .id(trace2.id())
                            .projectName(trace2.projectName())
                            .build())
                    .toList();

            var traceIdToScoresMap = Stream.concat(scores1.stream(), scores2.stream())
                    .collect(groupingBy(FeedbackScoreItem::id));

            // When storing the scores in batch, adding some more unrelated random ones
            var feedbackScoreBatch = factory.manufacturePojo(FeedbackScoreBatch.class);
            feedbackScoreBatch = feedbackScoreBatch.toBuilder()
                    .scores(Stream.concat(
                            feedbackScoreBatch.scores().stream(),
                            traceIdToScoresMap.values().stream().flatMap(List::stream))
                            .toList())
                    .build();

            createScoreAndAssert(feedbackScoreBatch, API_KEY, TEST_WORKSPACE);

            var experimentItems = IntStream.range(0, 10)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(List.of(experiment1, experiment2).get(i / 5).id())
                            .traceId(traces.get(i / 5).id())
                            .feedbackScores(traceIdToScoresMap.get(traces.get(i / 5).id()).stream()
                                    .map(FeedbackScoreMapper.INSTANCE::toFeedbackScore)
                                    .toList())
                            .build())
                    .toList();

            // When storing the experiment items in batch, adding some more unrelated random ones
            var experimentItemsBatch = factory.manufacturePojo(ExperimentItemsBatch.class);
            experimentItemsBatch = experimentItemsBatch.toBuilder()
                    .experimentItems(Stream.concat(
                            experimentItemsBatch.experimentItems().stream(),
                            experimentItems.stream())
                            .collect(toUnmodifiableSet()))
                    .build();

            Instant beforeCreateExperimentItems = Instant.now();

            createAndAssert(experimentItemsBatch, API_KEY, TEST_WORKSPACE);

            var actualDataset = getAndAssertEquals(dataset.id(), dataset, TEST_WORKSPACE, API_KEY);

            assertThat(actualDataset.experimentCount()).isEqualTo(2);
            assertThat(actualDataset.datasetItemsCount()).isEqualTo(datasetItems.size());
            assertThat(actualDataset.mostRecentExperimentAt()).isAfter(beforeCreateExperimentItems);
        }

    }

    private void createScoreAndAssert(FeedbackScoreBatchContainer feedbackScoreBatch, String apiKey,
            String workspaceName) {
        try (var actualResponse = client.target(getTracesPath())
                .path("feedback-scores")
                .request()
                .header(HttpHeaders.AUTHORIZATION, apiKey)
                .header(WORKSPACE_HEADER, workspaceName)
                .put(Entity.json(feedbackScoreBatch))) {

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
            assertThat(actualResponse.hasEntity()).isFalse();
        }
    }

    private void createAndAssert(Experiment experiment, String apiKey, String workspaceName) {
        experimentResourceClient.create(experiment, apiKey, workspaceName);
    }

    @Nested
    @DisplayName("Get:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class FindDatasets {

        @Test
        @DisplayName("Success")
        void getDatasets() {

            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> expected1 = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);
            List<Dataset> expected2 = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);

            expected1.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));
            expected2.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            Dataset dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .name("The most expressive LLM: " + UUID.randomUUID()
                            + " \uD83D\uDE05\uD83E\uDD23\uD83D\uDE02\uD83D\uDE42\uD83D\uDE43\uD83E\uDEE0")
                    .description("Emoji Test \uD83E\uDD13\uD83E\uDDD0")
                    .build();

            createAndAssert(dataset, apiKey, workspaceName);

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            int defaultPageSize = 10;

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);
            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);

            var expectedContent = new ArrayList<Dataset>();
            expectedContent.add(dataset);

            expected2.reversed()
                    .stream()
                    .filter(__ -> expectedContent.size() < defaultPageSize)
                    .forEach(expectedContent::add);

            expected1.reversed()
                    .stream()
                    .filter(__ -> expectedContent.size() < defaultPageSize)
                    .forEach(expectedContent::add);

            findAndAssertPage(actualEntity, defaultPageSize, expectedContent.size() + 1, 1, expectedContent);
        }

        @Test
        @DisplayName("when limit is 5 but there are N datasets, then return 5 datasets and total N")
        void getDatasets__whenLimitIs5ButThereAre10Datasets__thenReturn5DatasetsAndTotal10() {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> expected1 = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);
            List<Dataset> expected2 = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);

            expected1.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));
            expected2.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            int pageSize = 5;
            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", pageSize)
                    .queryParam("page", 1)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            var expectedContent = new ArrayList<>(expected2.reversed().subList(0, pageSize));

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);

            findAndAssertPage(actualEntity, pageSize, expected1.size() + expected2.size(), 1, expectedContent);
        }

        @Test
        @DisplayName("when fetching all datasets, then return datasets sorted by created date")
        void getDatasets__whenFetchingAllDatasets__thenReturnDatasetsSortedByCreatedDate() {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> expected = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);

            expected.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", 5)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);

            findAndAssertPage(actualEntity, expected.size(), expected.size(), 1, expected.reversed());
        }

        @ParameterizedTest
        @MethodSource("sortDirectionProvider")
        @DisplayName("when fetching all datasets, then return datasets sorted by last_created_experiment_at")
        void getDatasets__whenFetchingAllDatasets__thenReturnDatasetsSortedByLastCreatedExperimentAt(
                Direction requestDirection, Direction expectedDirection) {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> expected = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);
            Set<DatasetLastExperimentCreated> datasetsLastExperimentCreated = new HashSet<>();

            expected.forEach(dataset -> {
                var id = createAndAssert(dataset, apiKey, workspaceName);
                datasetsLastExperimentCreated.add(new DatasetLastExperimentCreated(id, Instant.now()));
            });

            saveDatasetsLastExperimentCreated(datasetsLastExperimentCreated, workspaceId);

            if (expectedDirection == Direction.DESC) {
                expected = expected.reversed();
            }

            requestAndAssertDatasetsPage(workspaceName, apiKey, expected, requestDirection,
                    SortableFields.LAST_CREATED_EXPERIMENT_AT, null);
        }

        @ParameterizedTest
        @MethodSource
        @DisplayName("when fetching all datasets, then return datasets sorted by valid fields")
        void getDatasets__whenFetchingAllDatasets__thenReturnDatasetsSortedByByValidFields(
                Comparator<Dataset> comparator,
                SortingField sorting) {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> datasets = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class).stream()
                    .map(d -> d.toBuilder()
                            .lastUpdatedBy(USER)
                            .createdBy(USER)
                            .build())
                    .toList();

            datasets.forEach(dataset -> {
                var id = createAndAssert(dataset, apiKey, workspaceName);

                saveDatasetsLastOptimizationCreated(Set.of(new DatasetLastOptimizationCreated(id, Instant.now())),
                        workspaceId);
                saveDatasetsLastExperimentCreated(Set.of(new DatasetLastExperimentCreated(id, Instant.now())),
                        workspaceId);
            });

            datasets = datasets.stream()
                    .sorted(comparator)
                    .toList();

            requestAndAssertDatasetsPage(workspaceName, apiKey, datasets, sorting.direction(),
                    sorting.field(), null);
        }

        private Stream<Arguments> getDatasets__whenFetchingAllDatasets__thenReturnDatasetsSortedByByValidFields() {
            // Comparators for all sortable fields
            Comparator<Dataset> idComparator = Comparator.comparing(Dataset::id);
            Comparator<Dataset> nameComparator = Comparator.comparing(Dataset::name, String.CASE_INSENSITIVE_ORDER);
            Comparator<Dataset> descriptionComparator = Comparator.comparing(Dataset::description,
                    String.CASE_INSENSITIVE_ORDER);
            Comparator<Dataset> tagsComparator = Comparator.comparing(d -> d.tags().toString(),
                    String.CASE_INSENSITIVE_ORDER);
            Comparator<Dataset> createdAtComparator = Comparator.comparing(Dataset::createdAt);
            Comparator<Dataset> createdByComparator = Comparator.comparing(Dataset::createdBy,
                    String.CASE_INSENSITIVE_ORDER);
            Comparator<Dataset> lastUpdatedAtComparator = Comparator.comparing(Dataset::lastUpdatedAt);
            Comparator<Dataset> lastUpdatedByComparator = Comparator.comparing(Dataset::lastUpdatedBy,
                    String.CASE_INSENSITIVE_ORDER);
            Comparator<Dataset> lastCreatedExperimentAtComparator = Comparator.comparing(
                    Dataset::lastCreatedExperimentAt,
                    Comparator.nullsLast(Comparator.naturalOrder()));
            Comparator<Dataset> lastCreatedOptimizationAtComparator = Comparator.comparing(
                    Dataset::lastCreatedOptimizationAt,
                    Comparator.nullsLast(Comparator.naturalOrder()));
            Comparator<Dataset> idComparatorReversed = Comparator.comparing(Dataset::id).reversed();

            return Stream.of(
                    // ID field sorting
                    Arguments.of(
                            idComparator,
                            SortingField.builder().field(SortableFields.ID).direction(Direction.ASC).build()),
                    Arguments.of(
                            idComparator.reversed(),
                            SortingField.builder().field(SortableFields.ID).direction(Direction.DESC).build()),

                    // NAME field sorting
                    Arguments.of(
                            nameComparator.thenComparing(idComparatorReversed),
                            SortingField.builder().field(SortableFields.NAME).direction(Direction.ASC).build()),
                    Arguments.of(
                            nameComparator.reversed().thenComparing(idComparatorReversed),
                            SortingField.builder().field(SortableFields.NAME).direction(Direction.DESC).build()),

                    // DESCRIPTION field sorting
                    Arguments.of(
                            descriptionComparator,
                            SortingField.builder().field(SortableFields.DESCRIPTION).direction(Direction.ASC).build()),
                    Arguments.of(
                            descriptionComparator.reversed(),
                            SortingField.builder().field(SortableFields.DESCRIPTION).direction(Direction.DESC).build()),

                    // TAGS field sorting
                    Arguments.of(
                            tagsComparator,
                            SortingField.builder().field(SortableFields.TAGS).direction(Direction.ASC).build()),
                    Arguments.of(
                            tagsComparator.reversed(),
                            SortingField.builder().field(SortableFields.TAGS).direction(Direction.DESC).build()),

                    // CREATED_AT field sorting
                    Arguments.of(
                            createdAtComparator,
                            SortingField.builder().field(SortableFields.CREATED_AT).direction(Direction.ASC).build()),
                    Arguments.of(
                            createdAtComparator.reversed(),
                            SortingField.builder().field(SortableFields.CREATED_AT).direction(Direction.DESC).build()),

                    // CREATED_BY field sorting
                    Arguments.of(
                            createdByComparator.thenComparing(idComparatorReversed),
                            SortingField.builder().field(SortableFields.CREATED_BY).direction(Direction.ASC).build()),
                    Arguments.of(
                            createdByComparator.reversed().thenComparing(idComparatorReversed),
                            SortingField.builder().field(SortableFields.CREATED_BY).direction(Direction.DESC).build()),

                    // LAST_UPDATED_AT field sorting
                    Arguments.of(
                            lastUpdatedAtComparator,
                            SortingField.builder().field(SortableFields.LAST_UPDATED_AT).direction(Direction.ASC)
                                    .build()),
                    Arguments.of(
                            lastUpdatedAtComparator.reversed(),
                            SortingField.builder().field(SortableFields.LAST_UPDATED_AT).direction(Direction.DESC)
                                    .build()),

                    // LAST_UPDATED_BY field sorting
                    Arguments.of(
                            lastUpdatedByComparator.thenComparing(idComparatorReversed),
                            SortingField.builder().field(SortableFields.LAST_UPDATED_BY).direction(Direction.ASC)
                                    .build()),
                    Arguments.of(
                            lastUpdatedByComparator.reversed().thenComparing(idComparatorReversed),
                            SortingField.builder().field(SortableFields.LAST_UPDATED_BY).direction(Direction.DESC)
                                    .build()),

                    // LAST_CREATED_EXPERIMENT_AT field sorting
                    Arguments.of(
                            lastCreatedExperimentAtComparator,
                            SortingField.builder().field(SortableFields.LAST_CREATED_EXPERIMENT_AT)
                                    .direction(Direction.ASC).build()),
                    Arguments.of(
                            lastCreatedExperimentAtComparator.reversed(),
                            SortingField.builder().field(SortableFields.LAST_CREATED_EXPERIMENT_AT)
                                    .direction(Direction.DESC).build()),

                    // LAST_CREATED_OPTIMIZATION_AT field sorting
                    Arguments.of(
                            lastCreatedOptimizationAtComparator,
                            SortingField.builder().field(SortableFields.LAST_CREATED_OPTIMIZATION_AT)
                                    .direction(Direction.ASC).build()),
                    Arguments.of(
                            lastCreatedOptimizationAtComparator.reversed(),
                            SortingField.builder().field(SortableFields.LAST_CREATED_OPTIMIZATION_AT)
                                    .direction(Direction.DESC).build()));
        }

        @ParameterizedTest
        @MethodSource("sortDirectionProvider")
        @DisplayName("when fetching all datasets, then return datasets sorted by last_created_optimization_at")
        void getDatasets__whenFetchingAllDatasets__thenReturnDatasetsSortedByLastCreatedOptimizationAt(
                Direction requestDirection, Direction expectedDirection) {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> expected = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);
            Set<DatasetLastOptimizationCreated> datasetsLastOptimizationCreated = new HashSet<>();

            expected.forEach(dataset -> {
                var id = createAndAssert(dataset, apiKey, workspaceName);
                datasetsLastOptimizationCreated.add(new DatasetLastOptimizationCreated(id, Instant.now()));
            });

            saveDatasetsLastOptimizationCreated(datasetsLastOptimizationCreated, workspaceId);

            if (expectedDirection == Direction.DESC) {
                expected = expected.reversed();
            }

            requestAndAssertDatasetsPage(workspaceName, apiKey, expected, requestDirection,
                    SortableFields.LAST_CREATED_OPTIMIZATION_AT, null);
        }

        private void saveDatasetsLastOptimizationCreated(
                Set<DatasetLastOptimizationCreated> datasetsLastOptimizationCreated, String workspaceId) {
            mySqlTemplate.inTransaction(WRITE, handle -> {

                var dao = handle.attach(DatasetDAO.class);
                dao.recordOptimizations(workspaceId, datasetsLastOptimizationCreated);

                return null;
            });
        }

        private void saveDatasetsLastExperimentCreated(Set<DatasetLastExperimentCreated> datasetsLastExperimentCreated,
                String workspaceId) {
            mySqlTemplate.inTransaction(WRITE, handle -> {

                var dao = handle.attach(DatasetDAO.class);
                dao.recordExperiments(workspaceId, datasetsLastExperimentCreated);

                return null;
            });
        }

        public static Stream<Arguments> sortDirectionProvider() {
            return Stream.of(
                    Arguments.of(Named.of("non specified", null), Direction.ASC),
                    Arguments.of(Named.of("ascending", Direction.ASC), Direction.ASC),
                    Arguments.of(Named.of("descending", Direction.DESC), Direction.DESC));
        }

        private void requestAndAssertDatasetsPage(String workspaceName, String apiKey, List<Dataset> allDatasets,
                Direction request, String sortingField, List<DatasetFilter> filters) {

            WebTarget target = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", Math.max(1, allDatasets.size()));

            if (sortingField != null) {
                var sorting = List.of(SortingField.builder()
                        .field(sortingField)
                        .direction(request)
                        .build());
                target = target.queryParam("sorting", URLEncoder.encode(JsonUtils.writeValueAsString(sorting),
                        StandardCharsets.UTF_8));
            }

            if (CollectionUtils.isNotEmpty(filters)) {
                target = target.queryParam("filters", toURLEncodedQueryParam(filters));
            }

            var actualResponse = target
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
            assertThat(actualEntity.size()).isEqualTo(allDatasets.size());
            assertThat(actualEntity.total()).isEqualTo(allDatasets.size());
            assertThat(actualEntity.page()).isEqualTo(1);

            assertThat(actualEntity.content())
                    .usingRecursiveFieldByFieldElementComparatorIgnoringFields(DATASET_IGNORED_FIELDS)
                    .containsExactlyElementsOf(allDatasets);
        }

        @ParameterizedTest
        @MethodSource("getValidFilters")
        @DisplayName("when fetching all datasets, then return datasets filtered datasets")
        void whenFilterDatasets__thenReturnDatasetsFiltered(Function<List<Dataset>, DatasetFilter> getFilter,
                Function<List<Dataset>, List<Dataset>> getExpectedDatasets) {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> datasets = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);

            Set<DatasetLastOptimizationCreated> datasetsLastOptimizationCreated = new HashSet<>();
            Set<DatasetLastExperimentCreated> datasetsLastExperimentCreated = new HashSet<>();

            datasets.forEach(dataset -> {
                var id = createAndAssert(dataset, apiKey, workspaceName);
                datasetsLastOptimizationCreated.add(new DatasetLastOptimizationCreated(id, Instant.now()));
                datasetsLastExperimentCreated.add(new DatasetLastExperimentCreated(id, Instant.now()));
            });

            saveDatasetsLastOptimizationCreated(datasetsLastOptimizationCreated, workspaceId);
            saveDatasetsLastExperimentCreated(datasetsLastExperimentCreated, workspaceId);

            List<Dataset> expectedDatasets = getExpectedDatasets.apply(datasets).reversed();
            DatasetFilter filter = getFilter.apply(datasets);

            requestAndAssertDatasetsPage(workspaceName, apiKey, expectedDatasets, null,
                    null, List.of(filter));
        }

        private Stream<Arguments> getValidFilters() {
            return Stream.of(
                    // TAGS field tests (existing)
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.TAGS)
                                    .operator(Operator.CONTAINS)
                                    .value(datasets.getFirst().tags().iterator().next())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> List.of(datasets.getFirst())),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.TAGS)
                                    .operator(Operator.NOT_CONTAINS)
                                    .value(datasets.getFirst().tags().iterator().next())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets.subList(1, datasets.size())),

                    // ID field tests
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.ID)
                                    .operator(Operator.EQUAL)
                                    .value(datasets.getFirst().id().toString())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> List.of(datasets.getFirst())),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.ID)
                                    .operator(Operator.NOT_EQUAL)
                                    .value(datasets.getFirst().id().toString())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets.subList(1, datasets.size())),

                    // NAME field tests
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.NAME)
                                    .operator(Operator.EQUAL)
                                    .value(datasets.getFirst().name())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> List.of(datasets.getFirst())),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.NAME)
                                    .operator(Operator.NOT_EQUAL)
                                    .value(datasets.getFirst().name())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets.subList(1, datasets.size())),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.NAME)
                                    .operator(Operator.CONTAINS)
                                    .value(datasets.getFirst().name().substring(0, 3))
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets.stream()
                                    .filter(dataset -> dataset.name()
                                            .contains(datasets.getFirst().name().substring(0, 3)))
                                    .toList()),

                    // DESCRIPTION field tests
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.DESCRIPTION)
                                    .operator(Operator.EQUAL)
                                    .value(datasets.getFirst().description())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> List.of(datasets.getFirst())),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.DESCRIPTION)
                                    .operator(Operator.NOT_EQUAL)
                                    .value(datasets.getFirst().description())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets.subList(1, datasets.size())),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.DESCRIPTION)
                                    .operator(Operator.CONTAINS)
                                    .value(datasets.getFirst().description() != null
                                            ? datasets.getFirst().description().substring(0,
                                                    Math.min(3, datasets.getFirst().description().length()))
                                            : "test")
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets.stream()
                                    .filter(dataset -> {
                                        String searchValue = datasets.getFirst().description() != null
                                                ? datasets.getFirst().description().substring(0,
                                                        Math.min(3, datasets.getFirst().description().length()))
                                                : "test";
                                        return dataset.description() != null
                                                && dataset.description().contains(searchValue);
                                    })
                                    .toList()),

                    // CREATED_AT field tests (following prompt test pattern)
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.CREATED_AT)
                                    .operator(Operator.NOT_EQUAL)
                                    .value(Instant.now().toString())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.CREATED_AT)
                                    .operator(Operator.GREATER_THAN)
                                    .value(Instant.now().minus(5, ChronoUnit.SECONDS).toString())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets),

                    // CREATED_BY field tests
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.CREATED_BY)
                                    .operator(Operator.EQUAL)
                                    .value(USER)
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.CREATED_BY)
                                    .operator(Operator.NOT_EQUAL)
                                    .value(USER)
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> List.of()),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.CREATED_BY)
                                    .operator(Operator.CONTAINS)
                                    .value(USER.substring(0, 3))
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets),

                    // LAST_UPDATED_AT field tests (following prompt test pattern)
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.LAST_UPDATED_AT)
                                    .operator(Operator.GREATER_THAN_EQUAL)
                                    .value(Instant.now().toString())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> List.of()),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.LAST_UPDATED_AT)
                                    .operator(Operator.LESS_THAN)
                                    .value(Instant.now().toString())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets),

                    // LAST_UPDATED_BY field tests
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.LAST_UPDATED_BY)
                                    .operator(Operator.EQUAL)
                                    .value(USER)
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.LAST_UPDATED_BY)
                                    .operator(Operator.NOT_EQUAL)
                                    .value(USER)
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> List.of()),
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.LAST_UPDATED_BY)
                                    .operator(Operator.CONTAINS)
                                    .value(USER.substring(0, 3))
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets),

                    // LAST_CREATED_EXPERIMENT_AT field tests
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.LAST_CREATED_EXPERIMENT_AT)
                                    .operator(Operator.LESS_THAN)
                                    .value(Instant.now().toString())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets),

                    // LAST_CREATED_OPTIMIZATION_AT field tests
                    Arguments.of(
                            (Function<List<Dataset>, DatasetFilter>) datasets -> DatasetFilter.builder()
                                    .field(DatasetField.LAST_CREATED_OPTIMIZATION_AT)
                                    .operator(Operator.LESS_THAN)
                                    .value(Instant.now().toString())
                                    .build(),
                            (Function<List<Dataset>, List<Dataset>>) datasets -> datasets));
        }

        @Test
        @DisplayName("when searching by dataset name, then return full text search result")
        void getDatasets__whenSearchingByDatasetName__thenReturnFullTextSearchResult() {
            UUID datasetSuffix = UUID.randomUUID();
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> datasets = List.of(
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("MySQL, realtime chatboot: " + datasetSuffix).build(),
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("Chatboot using mysql: " + datasetSuffix)
                            .build(),
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("Chatboot MYSQL expert: " + datasetSuffix)
                            .build(),
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("Chatboot expert (my SQL): " + datasetSuffix).build(),
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("Chatboot expert: " + datasetSuffix)
                            .build());

            datasets.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", 100)
                    .queryParam("name", "MySql")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
            assertThat(actualEntity.total()).isEqualTo(3);
            assertThat(actualEntity.size()).isEqualTo(3);

            var actualDatasets = actualEntity.content();
            assertThat(actualDatasets.stream().map(Dataset::name).toList()).contains(
                    "MySQL, realtime chatboot: " + datasetSuffix,
                    "Chatboot using mysql: " + datasetSuffix,
                    "Chatboot MYSQL expert: " + datasetSuffix);
        }

        @Test
        @DisplayName("when searching by dataset name fragments, then return full text search result")
        void getDatasets__whenSearchingByDatasetNameFragments__thenReturnFullTextSearchResult() {

            UUID datasetSuffix = UUID.randomUUID();
            String workspaceName = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> datasets = List.of(
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("MySQL: " + datasetSuffix)
                            .build(),
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("Chat-boot using mysql: " + datasetSuffix)
                            .build(),
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("MYSQL CHATBOOT expert: " + datasetSuffix)
                            .build(),
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("Expert Chatboot: " + datasetSuffix)
                            .build(),
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name("My chat expert: " + datasetSuffix)
                            .build());

            datasets.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", 100)
                    .queryParam("name", "cha")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
            assertThat(actualEntity.total()).isEqualTo(4);
            assertThat(actualEntity.size()).isEqualTo(4);

            var actualDatasets = actualEntity.content();

            assertThat(actualDatasets.stream().map(Dataset::name).toList()).contains(
                    "Chat-boot using mysql: " + datasetSuffix,
                    "MYSQL CHATBOOT expert: " + datasetSuffix,
                    "Expert Chatboot: " + datasetSuffix,
                    "My chat expert: " + datasetSuffix);
        }

        @Test
        @DisplayName("when searching by dataset name and workspace name, then return full text search result")
        void getDatasets__whenSearchingByDatasetNameAndWorkspaceName__thenReturnFullTextSearchResult() {

            var name = UUID.randomUUID().toString();
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            Dataset dataset1 = factory.manufacturePojo(Dataset.class).toBuilder()
                    .name(name)
                    .build();

            Dataset dataset2 = factory.manufacturePojo(Dataset.class).toBuilder()
                    .name(name)
                    .build();

            createAndAssert(dataset1, API_KEY, TEST_WORKSPACE);

            createAndAssert(dataset2, apiKey, workspaceName);

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", 100)
                    .queryParam("name", name)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            findAndAssertPage(actualEntity, 1, 1, 1, List.of(dataset1));
        }

        @Test
        @DisplayName("when searching by dataset name with different workspace name, then return no match")
        void getDatasets__whenSearchingByDatasetNameWithDifferentWorkspaceAndWorkspaceName__thenReturnNoMatch() {

            var name = UUID.randomUUID().toString();
            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> datasets = List.of(
                    factory.manufacturePojo(Dataset.class).toBuilder()
                            .name(name)
                            .build());

            datasets.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", 100)
                    .queryParam("name", name)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            findAndAssertPage(actualEntity, 0, 0, 1, List.of());
        }

        @Test
        @DisplayName("when datasets have experiments linked to them, then return datasets with experiment summary")
        void getDatasets__whenDatasetsHaveExperimentsLinkedToThem__thenReturnDatasetsWithExperimentSummary() {

            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> datasets = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);

            datasets.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            AtomicInteger index = new AtomicInteger();

            var experiments = experimentResourceClient.generateExperimentList()
                    .stream()
                    .flatMap(experiment -> Stream.of(experiment.toBuilder()
                            .datasetName(datasets.get(index.getAndIncrement()).name())
                            .datasetId(null)
                            .build()))
                    .toList();

            experiments.forEach(experiment -> createAndAssert(experiment, apiKey, workspaceName));

            index.set(0);

            var datasetItems = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class);

            datasetItems.forEach(datasetItem -> putAndAssert(
                    new DatasetItemBatch(null, datasets.get(index.getAndIncrement()).id(), List.of(datasetItem)),
                    workspaceName, apiKey));

            // Creating two traces with input, output and scores
            var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class);

            traces.forEach(trace -> createTrace(trace, apiKey, workspaceName));

            index.set(0);

            // Creating 5 scores peach each of the two traces above
            var scores = PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItem.class)
                    .stream()
                    .map(feedbackScoreBatchItem -> feedbackScoreBatchItem.toBuilder()
                            .id(traces.get(index.get()).id())
                            .projectName(traces.get(index.getAndIncrement()).projectName())
                            .build())
                    .toList();

            var traceIdToScoresMap = scores.stream()
                    .collect(groupingBy(FeedbackScoreItem::id));

            // When storing the scores in batch, adding some more unrelated random ones
            var feedbackScoreBatch = factory.manufacturePojo(FeedbackScoreBatch.class);
            feedbackScoreBatch = feedbackScoreBatch.toBuilder()
                    .scores(Stream.concat(
                            feedbackScoreBatch.scores().stream(),
                            traceIdToScoresMap.values().stream().flatMap(List::stream))
                            .toList())
                    .build();

            createScoreAndAssert(feedbackScoreBatch, apiKey, workspaceName);

            index.set(0);

            var experimentItems = IntStream.range(0, 5)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .datasetItemId(datasetItems.get(index.get()).id())
                            .experimentId(experiments.get(index.get()).id())
                            .traceId(traces.get(index.get()).id())
                            .feedbackScores(traceIdToScoresMap.get(traces.get(index.getAndIncrement()).id()).stream()
                                    .map(FeedbackScoreMapper.INSTANCE::toFeedbackScore)
                                    .toList())
                            .build())
                    .collect(toSet());

            var experimentItemsBatch = factory.manufacturePojo(ExperimentItemsBatch.class).toBuilder()
                    .experimentItems(experimentItems)
                    .build();
            Instant beforeCreateExperimentItems = Instant.now();

            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);

            assertThat(actualEntity.content()).hasSize(datasets.size());
            assertThat(actualEntity.total()).isEqualTo(datasets.size());
            assertThat(actualEntity.page()).isEqualTo(1);
            assertThat(actualEntity.size()).isEqualTo(datasets.size());

            for (int i = 0; i < actualEntity.content().size(); i++) {
                var dataset = actualEntity.content().get(i);

                assertThat(dataset.experimentCount()).isEqualTo(1);
                assertThat(dataset.datasetItemsCount()).isEqualTo(1);
                assertThat(dataset.mostRecentExperimentAt()).isAfter(beforeCreateExperimentItems);
            }
        }

        @Test
        @DisplayName("when searching by dataset with experiments only but no experiment found, then return empty page")
        void getDatasets__whenSearchingByDatasetWithExperimentsOnlyButNoExperimentFound__thenReturnEmptyPage() {

            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> datasets = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class);

            datasets.forEach(dataset -> createAndAssert(dataset, apiKey, workspaceName));

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("with_experiments_only", true)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            findAndAssertPage(actualEntity, 0, 0, 1, List.of());
        }

        @ParameterizedTest
        @MethodSource
        @DisplayName("when searching by dataset with experiments only and result having {} datasets, then return page")
        void getDatasets__whenSearchingByDatasetWithExperimentsOnlyAndResultHavingXDatasets__thenReturnPage(
                int datasetCount) {

            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> expectedDatasets = IntStream.range(0, datasetCount)
                    .parallel()
                    .mapToObj(i -> createDatasetWithExperiment(apiKey, workspaceName, null))
                    .sorted(Comparator.comparing(Dataset::id))
                    .toList();

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", expectedDatasets.size())
                    .queryParam("with_experiments_only", true)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            findAndAssertPage(actualEntity, expectedDatasets.size(), expectedDatasets.size(), 1,
                    expectedDatasets.reversed());
        }

        Stream<Arguments> getDatasets__whenSearchingByDatasetWithExperimentsOnlyAndResultHavingXDatasets__thenReturnPage() {
            return Stream.of(
                    arguments(10),
                    arguments(100),
                    arguments(110));
        }

        @Test
        @DisplayName("when searching by dataset with optimizations only and result having {} datasets, then return page")
        void getDatasets__whenSearchingByDatasetWithOptimizationsOnlyAndResultHavingXDatasets__thenReturnPage() {

            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            List<Dataset> expectedDatasets = IntStream.range(0, 10)
                    .parallel()
                    .mapToObj(i -> createDatasetWithOptimization(apiKey, workspaceName))
                    .sorted(Comparator.comparing(Dataset::id))
                    .toList();

            Awaitility.await()
                    .atMost(Duration.ofSeconds(10))
                    .untilAsserted(() -> {
                        var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                                .queryParam("size", expectedDatasets.size())
                                .queryParam("with_optimizations_only", true)
                                .request()
                                .header(HttpHeaders.AUTHORIZATION, apiKey)
                                .header(WORKSPACE_HEADER, workspaceName)
                                .get();

                        var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

                        findAndAssertPage(actualEntity, expectedDatasets.size(), expectedDatasets.size(), 1,
                                expectedDatasets.reversed());
                    });

        }

        @ParameterizedTest
        @MethodSource
        @DisplayName("when searching by dataset with experiments only, name {}, and result having {} datasets, then return page")
        void getDatasets__whenSearchingByDatasetWithExperimentsOnlyAndNameXAndResultHavingXDatasets__thenReturnPage(
                String datasetNamePrefix, int datasetCount, int expectedMatchCount) {

            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            int unexpectedDatasetCount = datasetCount - expectedMatchCount;

            IntStream.range(0, unexpectedDatasetCount)
                    .parallel()
                    .mapToObj(i -> createDatasetWithExperiment(apiKey, workspaceName, null))
                    .sorted(Comparator.comparing(Dataset::id))
                    .toList();

            List<Dataset> expectedMatchedDatasets = IntStream.range(0, expectedMatchCount)
                    .parallel()
                    .mapToObj(i -> createDatasetWithExperiment(apiKey, workspaceName, datasetNamePrefix))
                    .sorted(Comparator.comparing(Dataset::id))
                    .toList();

            Awaitility.await()
                    .atMost(Duration.ofSeconds(10))
                    .untilAsserted(() -> {
                        var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                                .queryParam("size", expectedMatchedDatasets.size())
                                .queryParam("with_experiments_only", true)
                                .queryParam("name", datasetNamePrefix)
                                .request()
                                .header(HttpHeaders.AUTHORIZATION, apiKey)
                                .header(WORKSPACE_HEADER, workspaceName)
                                .get();

                        var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

                        findAndAssertPage(actualEntity, expectedMatchedDatasets.size(), expectedMatchedDatasets.size(),
                                1,
                                expectedMatchedDatasets.reversed());
                    });
        }

        Stream<Arguments> getDatasets__whenSearchingByDatasetWithExperimentsOnlyAndNameXAndResultHavingXDatasets__thenReturnPage() {
            return Stream.of(
                    arguments(UUID.randomUUID().toString(), 10, 5),
                    arguments(UUID.randomUUID().toString(), 100, 50),
                    arguments(UUID.randomUUID().toString(), 110, 10));
        }

        private Dataset createDatasetWithExperiment(String apiKey, String workspaceName, String datasetNamePrefix) {
            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .name(datasetNamePrefix == null
                            ? UUID.randomUUID().toString()
                            : datasetNamePrefix + " " + UUID.randomUUID())
                    .build();

            createAndAssert(dataset, apiKey, workspaceName);

            Experiment experiment = experimentResourceClient.createPartialExperiment()
                    .datasetName(dataset.name())
                    .type(null)
                    .build();

            createAndAssert(
                    experiment,
                    apiKey,
                    workspaceName);

            experiment = getExperiment(apiKey, workspaceName, experiment);

            return dataset.toBuilder()
                    .experimentCount(1L)
                    .lastCreatedExperimentAt(experiment.createdAt())
                    .mostRecentExperimentAt(experiment.createdAt())
                    .build();
        }

        private Dataset createDatasetWithOptimization(String apiKey, String workspaceName) {
            var dataset = factory.manufacturePojo(Dataset.class);
            createAndAssert(dataset, apiKey, workspaceName);

            var optimization = optimizationResourceClient.createPartialOptimization().datasetName(dataset.name())
                    .build();
            optimizationResourceClient.create(optimization, apiKey, workspaceName);

            return dataset.toBuilder()
                    .optimizationCount(1L)
                    .build();
        }

        @ParameterizedTest
        @MethodSource
        @DisplayName("when searching by prompt id and result having {} datasets linked to experiments with prompt id, then return page")
        void getDatasets__whenSearchingByPromptIdAndResultHavingXDatasetsLinkedToExperimentsWithPromptId__thenReturnPage(
                int datasetCount, int expectedMatchCount) {

            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            IntStream.range(0, datasetCount - expectedMatchCount)
                    .parallel()
                    .forEach(i -> createDatasetWithExperiment(apiKey, workspaceName, null));

            Prompt prompt = factory.manufacturePojo(Prompt.class).toBuilder().latestVersion(null).build();

            PromptVersion promptVersion = promptResourceClient.createPromptVersion(prompt, apiKey, workspaceName);

            List<Dataset> expectedDatasets = IntStream.range(0, expectedMatchCount)
                    .parallel()
                    .mapToObj(i -> {
                        var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                                .name(UUID.randomUUID().toString())
                                .build();

                        createAndAssert(dataset, apiKey, workspaceName);

                        Experiment experiment = factory.manufacturePojo(Experiment.class).toBuilder()
                                .datasetName(dataset.name())
                                .type(null)
                                .promptVersion(
                                        Experiment.PromptVersionLink.builder()
                                                .promptId(promptVersion.promptId())
                                                .id(promptVersion.id())
                                                .commit(promptVersion.commit())
                                                .build())
                                .promptVersions(null)
                                .build();

                        createAndAssert(
                                experiment,
                                apiKey,
                                workspaceName);

                        experiment = getExperiment(apiKey, workspaceName, experiment);

                        // Create trial experiment for the same dataset, should not be included in experiment count
                        Experiment trial = factory.manufacturePojo(Experiment.class).toBuilder()
                                .datasetName(dataset.name())
                                .type(ExperimentType.TRIAL)
                                .promptVersion(null)
                                .promptVersions(null)
                                .build();

                        createAndAssert(
                                trial,
                                apiKey,
                                workspaceName);

                        return dataset.toBuilder()
                                .experimentCount(1L)
                                .lastCreatedExperimentAt(experiment.createdAt())
                                .mostRecentExperimentAt(experiment.createdAt())
                                .build();
                    })
                    .sorted(Comparator.comparing(Dataset::id))
                    .toList();

            var actualEntity = datasetResourceClient.getDatasetPage(apiKey, workspaceName, expectedDatasets.size(),
                    promptVersion);

            findAndAssertPage(actualEntity, expectedDatasets.size(), expectedDatasets.size(), 1,
                    expectedDatasets.reversed());
        }

        @ParameterizedTest
        @MethodSource("getDatasets__whenSearchingByPromptIdAndResultHavingXDatasetsLinkedToExperimentsWithPromptId__thenReturnPage")
        @DisplayName("when searching by prompt id and result having {} datasets linked to experiments with list of prompt ids, then return page")
        void getDatasets__whenSearchingByPromptIdAndResultHavingXDatasetsLinkedToExperimentsWithListOfPromptIds__thenReturnPage(
                int datasetCount, int expectedMatchCount) {

            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            IntStream.range(0, datasetCount - expectedMatchCount)
                    .parallel()
                    .forEach(i -> createDatasetWithExperiment(apiKey, workspaceName, null));

            Prompt prompt = factory.manufacturePojo(Prompt.class).toBuilder().latestVersion(null).build();

            PromptVersion promptVersion = promptResourceClient.createPromptVersion(prompt, apiKey, workspaceName);

            List<Dataset> expectedDatasets = IntStream.range(0, expectedMatchCount)
                    .parallel()
                    .mapToObj(i -> {
                        var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                                .name(UUID.randomUUID().toString())
                                .build();

                        createAndAssert(dataset, apiKey, workspaceName);

                        Experiment experiment = factory.manufacturePojo(Experiment.class).toBuilder()
                                .datasetName(dataset.name())
                                .type(null)
                                .promptVersion(null)
                                .promptVersions(List.of(
                                        Experiment.PromptVersionLink.builder()
                                                .promptId(promptVersion.promptId())
                                                .id(promptVersion.id())
                                                .commit(promptVersion.commit())
                                                .build()))
                                .build();

                        createAndAssert(
                                experiment,
                                apiKey,
                                workspaceName);

                        experiment = getExperiment(apiKey, workspaceName, experiment);

                        return dataset.toBuilder()
                                .experimentCount(1L)
                                .lastCreatedExperimentAt(experiment.createdAt())
                                .mostRecentExperimentAt(experiment.createdAt())
                                .build();
                    })
                    .sorted(Comparator.comparing(Dataset::id))
                    .toList();

            var actualEntity = datasetResourceClient.getDatasetPage(apiKey, workspaceName, expectedDatasets.size(),
                    promptVersion);

            findAndAssertPage(actualEntity, expectedDatasets.size(), expectedDatasets.size(), 1,
                    expectedDatasets.reversed());
        }

        Stream<Arguments> getDatasets__whenSearchingByPromptIdAndResultHavingXDatasetsLinkedToExperimentsWithPromptId__thenReturnPage() {
            return Stream.of(
                    arguments(10, 0),
                    arguments(10, 5));
        }
    }

    private Experiment getExperiment(String apiKey, String workspaceName, Experiment experiment) {

        try (var actualResponse = client.target(EXPERIMENT_RESOURCE_URI.formatted(baseURI))
                .path(experiment.id().toString())
                .request()
                .header(HttpHeaders.AUTHORIZATION, apiKey)
                .header(WORKSPACE_HEADER, workspaceName)
                .get()) {

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
            assertThat(actualResponse.hasEntity()).isTrue();

            return actualResponse.readEntity(Experiment.class);
        }

    }

    private void findAndAssertPage(Dataset.DatasetPage actualEntity, int expected, int total, int page,
            List<Dataset> expectedContent) {
        assertThat(actualEntity.size()).isEqualTo(expected);
        assertThat(actualEntity.content()).hasSize(expected);
        assertThat(actualEntity.page()).isEqualTo(page);
        assertThat(actualEntity.total()).isEqualTo(total);

        assertThat(actualEntity.content())
                .usingRecursiveFieldByFieldElementComparatorIgnoringFields(DATASET_IGNORED_FIELDS)
                .isEqualTo(expectedContent);
    }

    @Nested
    @DisplayName("Update:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class UpdateDataset {

        public Stream<Arguments> invalidDataset() {
            return Stream.of(
                    arguments(factory.manufacturePojo(DatasetUpdate.class).toBuilder().name(null).build(),
                            "name must not be blank"),
                    arguments(factory.manufacturePojo(DatasetUpdate.class).toBuilder().name("").build(),
                            "name must not be blank"),
                    arguments(
                            factory.manufacturePojo(DatasetUpdate.class).toBuilder().description("").build(),
                            "description must not be blank"),
                    arguments(
                            factory.manufacturePojo(DatasetUpdate.class).toBuilder().description("a".repeat(256))
                                    .build(),
                            "description cannot exceed 255 characters"));
        }

        @ParameterizedTest
        @MethodSource
        @DisplayName("Success")
        void updateDataset(DatasetUpdate datasetUpdate) {
            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            var id = createAndAssert(dataset);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.entity(datasetUpdate, MediaType.APPLICATION_JSON_TYPE))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                assertThat(actualResponse.hasEntity()).isFalse();
            }

            var expectedDataset = dataset.toBuilder()
                    .name(datasetUpdate.name())
                    .description(datasetUpdate.description())
                    .visibility(datasetUpdate.visibility())
                    .tags(datasetUpdate.tags() == null ? dataset.tags() : datasetUpdate.tags())
                    .build();

            getAndAssertEquals(id, expectedDataset, TEST_WORKSPACE, API_KEY);
        }

        Stream<Arguments> updateDataset() {
            return Stream.of(
                    arguments(factory.manufacturePojo(DatasetUpdate.class)),
                    arguments(factory.manufacturePojo(DatasetUpdate.class).toBuilder().tags(Set.of()).build()),
                    arguments(factory.manufacturePojo(DatasetUpdate.class).toBuilder().tags(null).build()));
        }

        @Test
        @DisplayName("when dataset not found, then return 404")
        void updateDataset__whenDatasetNotFound__thenReturn404() {
            var datasetUpdate = factory.manufacturePojo(DatasetUpdate.class);
            var id = UUID.randomUUID().toString();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id)
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.entity(datasetUpdate, MediaType.APPLICATION_JSON_TYPE))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
            }
        }

        @ParameterizedTest
        @MethodSource("invalidDataset")
        @DisplayName("when updating request is not valid, then return 422")
        void updateDataset__whenUpdatingRequestIsNotValid__thenReturn422(DatasetUpdate datasetUpdate,
                String errorMessage) {
            var id = factory.manufacturePojo(Dataset.class).id();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.entity(datasetUpdate, MediaType.APPLICATION_JSON_TYPE))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(errorMessage);
            }

        }

        @Test
        @DisplayName("when updating only name, then update only name")
        void updateDataset__whenUpdatingOnlyName__thenUpdateOnlyName() {
            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            var id = createAndAssert(dataset);

            var datasetUpdate = factory.manufacturePojo(DatasetUpdate.class)
                    .toBuilder()
                    .description(null)
                    .visibility(null)
                    .tags(null)
                    .build();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.entity(datasetUpdate, MediaType.APPLICATION_JSON_TYPE))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                assertThat(actualResponse.hasEntity()).isFalse();
            }

            var expectedDataset = dataset.toBuilder().name(datasetUpdate.name())
                    .description(datasetUpdate.description()).build();
            getAndAssertEquals(id, expectedDataset, TEST_WORKSPACE, API_KEY);
        }
    }

    @Nested
    @DisplayName("Delete:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class DeleteDataset {

        @Test
        @DisplayName("Success")
        void deleteDataset() {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            var id = createAndAssert(dataset);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .delete()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                assertThat(actualResponse.hasEntity()).isFalse();
            }

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get();

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
            assertThat(actualResponse.hasEntity()).isTrue();
            assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
        }

        @Test
        @DisplayName("delete batch datasets")
        void deleteDatasetsBatch() {
            var apiKey = UUID.randomUUID().toString();
            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var ids = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class).stream()
                    .map(dataset -> createAndAssert(dataset, apiKey, workspaceName)).toList();
            var idsToDelete = ids.subList(0, 3);
            var notDeletedIds = ids.subList(3, ids.size());

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("delete-batch")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .post(Entity.json(new BatchDelete(new HashSet<>(idsToDelete))))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_NO_CONTENT);
                assertThat(actualResponse.hasEntity()).isFalse();
            }

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .queryParam("size", ids.size())
                    .queryParam("page", 1)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get();

            var actualEntity = actualResponse.readEntity(Dataset.DatasetPage.class);

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_OK);
            assertThat(actualEntity.size()).isEqualTo(notDeletedIds.size());
            assertThat(actualEntity.content().stream().map(Dataset::id).toList())
                    .usingRecursiveComparison()
                    .ignoringCollectionOrder()
                    .isEqualTo(notDeletedIds);
        }

        @Test
        @DisplayName("when dataset does not exists, then return no content")
        void deleteDataset__whenDatasetDoesNotExists__thenReturnNoContent() {
            var id = UUID.randomUUID().toString();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id)
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .delete()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                assertThat(actualResponse.hasEntity()).isFalse();
            }
        }

        @Test
        @DisplayName("when deleting by dataset name, then return no content")
        void deleteDataset__whenDeletingByDatasetName__thenReturnNoContent() {
            var dataset = factory.manufacturePojo(Dataset.class);

            var id = createAndAssert(dataset);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("delete")
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(new DatasetIdentifier(dataset.name())))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                assertThat(actualResponse.hasEntity()).isFalse();
            }

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(id.toString())
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get();

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
            assertThat(actualResponse.hasEntity()).isTrue();
            assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
        }

        @Test
        @DisplayName("when deleting by dataset name and dataset does not exist, then return no content")
        void deleteDataset__whenDeletingByDatasetNameAndDatasetDoesNotExist__thenReturnNotFound() {
            var dataset = factory.manufacturePojo(Dataset.class);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("delete")
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(new DatasetIdentifier(dataset.name())))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
                assertThat(actualResponse.hasEntity()).isTrue();
            }

        }

        @ParameterizedTest
        @MethodSource
        @DisplayName("when deleting dataset should update optimization dataset_deleted")
        void deletingDataset__shouldUpdateOptimizationDatasetDeleted__thenReturnNotFound(
                Consumer<Dataset> datasetDeleteAction) {
            var dataset = factory.manufacturePojo(Dataset.class);
            createAndAssert(dataset);

            var optimizations = IntStream.range(0, 10)
                    .mapToObj(i -> {
                        var optimization = optimizationResourceClient.createPartialOptimization()
                                .datasetName(dataset.name())
                                .build();
                        var optimizationId = optimizationResourceClient.create(optimization, API_KEY, TEST_WORKSPACE);

                        return optimization.toBuilder().id(optimizationId).datasetName(null).build();
                    })
                    .toList();

            // Check that we do not have optimizations with deleted datasets
            var page = optimizationResourceClient.find(API_KEY, TEST_WORKSPACE, 1, optimizations.size(), null, null,
                    true, 200);
            assertThat(page.size()).isEqualTo(0);

            // Delete dataset
            datasetDeleteAction.accept(dataset);

            // Check that now we have one optimization with deleted datasets
            Awaitility.await()
                    .atMost(Duration.ofSeconds(10))
                    .untilAsserted(() -> {
                        var pageWithDeleted = optimizationResourceClient.find(API_KEY, TEST_WORKSPACE, 1,
                                optimizations.size(), null,
                                null, true, 200);
                        assertThat(pageWithDeleted.size()).isEqualTo(optimizations.size());

                        assertThat(pageWithDeleted.content())
                                .usingRecursiveComparison()
                                .ignoringFields(OPTIMIZATION_IGNORED_FIELDS)
                                .withComparatorForType(StatsUtils::bigDecimalComparator, BigDecimal.class)
                                .isEqualTo(optimizations.reversed());
                    });

            // Clean up
            optimizationResourceClient.delete(optimizations.stream().map(Optimization::id).collect(toSet()), API_KEY,
                    TEST_WORKSPACE);
        }

        private Stream<Consumer<Dataset>> deletingDataset__shouldUpdateOptimizationDatasetDeleted__thenReturnNotFound() {
            return Stream.of(
                    dataset -> datasetResourceClient.deleteDataset(dataset.id(), API_KEY, TEST_WORKSPACE),
                    dataset -> datasetResourceClient.deleteDatasetByName(dataset.name(), API_KEY, TEST_WORKSPACE),
                    dataset -> datasetResourceClient.deleteDatasetsBatch(Set.of(dataset.id()), API_KEY,
                            TEST_WORKSPACE));
        }
    }

    @Nested
    @DisplayName("Create dataset items:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class CreateDatasetItems {

        @Test
        @DisplayName("Success")
        void createDatasetItem() {
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .id(null)
                    .build();

            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .id(null)
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);
        }

        @Test
        @DisplayName("when item id is null, then return no content and create item")
        void createDatasetItem__whenItemIdIsNull__thenReturnNoContentAndCreateItem() {
            var item = factory.manufacturePojo(DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            getItemAndAssert(item, TEST_WORKSPACE, API_KEY);
        }

        @ParameterizedTest
        @MethodSource("invalidDatasetItemBatches")
        @DisplayName("when dataset item batch is not valid, then return 422")
        void createDatasetItem__whenDatasetItemIsNotValid__thenReturn422(DatasetItemBatch batch, String errorMessage) {
            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.json(batch))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(errorMessage);
            }
        }

        public Stream<Arguments> invalidDatasetItemBatches() {
            return Stream.of(
                    arguments(
                            factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                                    .items(List.of()).build(),
                            "items size must be between 1 and 1000"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(null).build(),
                            "items must not be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .datasetName(null)
                            .datasetId(null)
                            .build(),
                            "The request body must provide either a dataset_name or a dataset_id"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .datasetName("")
                            .datasetId(null)
                            .build(),
                            "datasetName must not be blank"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .datasetId(null)
                            .items(IntStream.range(0, 1001).mapToObj(i -> factory.manufacturePojo(DatasetItem.class))
                                    .toList())
                            .build(),
                            "items size must be between 1 and 1000"));
        }

        @Test
        @DisplayName("when dataset id not found, then return 404")
        void createDatasetItem__whenDatasetIdNotFound__thenReturn404() {

            var batch = factory.manufacturePojo(DatasetItemBatch.class);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.entity(batch, MediaType.APPLICATION_JSON_TYPE))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset not found");
            }
        }

        @Test
        @DisplayName("when dataset item id not valid, then return bad request")
        void createDatasetItem__whenDatasetItemIdIsNotValid__thenReturnBadRequest() {

            var item = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .id(UUID.randomUUID())
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(null)
                    .build();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.json(batch))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(400);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors())
                        .contains("dataset_item id must be a version 7 UUID");
            }
        }

        @Test
        @DisplayName("when dataset item already exists, then return no content and update item")
        void createDatasetItem__whenDatasetItemAlreadyExists__thenReturnNoContentAndUpdateItem() {
            var item = factory.manufacturePojo(DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            getItemAndAssert(item, TEST_WORKSPACE, API_KEY);

            var newItem = factory.manufacturePojo(DatasetItem.class)
                    .toBuilder()
                    .id(item.id())
                    .build();

            putAndAssert(batch.toBuilder()
                    .items(List.of(newItem))
                    .build(), TEST_WORKSPACE, API_KEY);

            getItemAndAssert(newItem, TEST_WORKSPACE, API_KEY);
        }

        @Test
        @DisplayName("when dataset item support null values for data fields, then return no content and create item")
        void createDatasetItem__whenDatasetItemSupportNullValuesForDataFields__thenReturnNoContentAndCreateItem() {
            var item = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .data(Map.of("test", NullNode.getInstance()))
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            getItemAndAssert(item, TEST_WORKSPACE, API_KEY);
        }

        @ParameterizedTest
        @MethodSource("invalidDatasetItems")
        @DisplayName("when dataset item batch contains duplicate items, then return 422")
        void createDatasetItem__whenDatasetItemBatchContainsDuplicateItems__thenReturn422(DatasetItemBatch batch,
                String errorMessage) {

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .put(Entity.entity(batch, MediaType.APPLICATION_JSON_TYPE))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(errorMessage);
            }
        }

        @Test
        @DisplayName("when dataset multiple items, then return no content and create items")
        void createDatasetItem__whenDatasetMultipleItems__thenReturnNoContentAndCreateItems() {
            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            items.forEach(item -> DatasetsResourceTest.this.getItemAndAssert(item, TEST_WORKSPACE, API_KEY));
        }

        public Stream<Arguments> invalidDatasetItems() {
            return Stream.of(
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .data(null)
                                    .build()))
                            .build(),
                            "items[0].data must not be empty"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .data(Map.of())
                                    .build()))
                            .build(),
                            "items[0].data must not be empty"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(null)
                                    .build()))
                            .build(),
                            "items[0].source must not be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(DatasetItemSource.MANUAL)
                                    .spanId(factory.manufacturePojo(UUID.class))
                                    .traceId(null)
                                    .build()))
                            .build(),
                            "items[0].source when it is manual, span_id must be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(DatasetItemSource.MANUAL)
                                    .spanId(null)
                                    .traceId(factory.manufacturePojo(UUID.class))
                                    .build()))
                            .build(),
                            "items[0].source when it is manual, trace_id must be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(DatasetItemSource.SDK)
                                    .spanId(factory.manufacturePojo(UUID.class))
                                    .traceId(null)
                                    .build()))
                            .build(),
                            "items[0].source when it is sdk, span_id must be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(DatasetItemSource.SDK)
                                    .traceId(factory.manufacturePojo(UUID.class))
                                    .spanId(null)
                                    .build()))
                            .build(),
                            "items[0].source when it is sdk, trace_id must be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(DatasetItemSource.SPAN)
                                    .spanId(null)
                                    .traceId(factory.manufacturePojo(UUID.class))
                                    .build()))
                            .build(),
                            "items[0].source when it is span, span_id must not be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(DatasetItemSource.SPAN)
                                    .traceId(null)
                                    .spanId(factory.manufacturePojo(UUID.class))
                                    .build()))
                            .build(),
                            "items[0].source when it is span, trace_id must not be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(DatasetItemSource.TRACE)
                                    .spanId(factory.manufacturePojo(UUID.class))
                                    .traceId(factory.manufacturePojo(UUID.class))
                                    .build()))
                            .build(),
                            "items[0].source when it is trace, span_id must be null"),
                    arguments(factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                            .items(List.of(factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .source(DatasetItemSource.TRACE)
                                    .spanId(null)
                                    .traceId(null)
                                    .build()))
                            .build(),
                            "items[0].source when it is trace, trace_id must not be null"));
        }

        @Test
        @DisplayName("when dataset item batch has max size, then return no content and create")
        void createDatasetItem__whenDatasetItemBatchHasMaxSize__thenReturnNoContentAndCreate() {

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build();

            UUID id = createAndAssert(dataset);

            var items = IntStream.range(0, 1000)
                    .mapToObj(__ -> factory.manufacturePojo(DatasetItem.class).toBuilder()
                            .experimentItems(null)
                            .createdAt(null)
                            .lastUpdatedAt(null)
                            .build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(id)
                    .datasetName(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);
        }

        @Test
        @DisplayName("when dataset item workspace and trace workspace does not match, then return conflict")
        void createDatasetItem__whenDatasetItemWorkspaceAndTraceWorkspaceDoesNotMatch__thenReturnConflict() {

            String workspaceName2 = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName2, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);

            var datasetId = createAndAssert(dataset, apiKey, workspaceName2);

            UUID traceId = createTrace(factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(UUID.randomUUID().toString())
                    .build(), API_KEY, TEST_WORKSPACE);

            var item = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .traceId(traceId)
                    .spanId(null)
                    .source(DatasetItemSource.TRACE)
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(datasetId)
                    .build();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName2)
                    .put(Entity.entity(batch, MediaType.APPLICATION_JSON_TYPE))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(409);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(
                        "trace workspace and dataset item workspace does not match");
            }
        }

        @Test
        @DisplayName("when dataset item workspace and span workspace does not match, then return conflict")
        void createDatasetItem__whenDatasetItemWorkspaceAndSpanWorkspaceDoesNotMatch__thenReturnConflict() {

            String workspaceName1 = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            String projectName = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName1, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);

            var datasetId = createAndAssert(dataset, apiKey, workspaceName1);

            UUID traceId = createTrace(factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(projectName)
                    .build(), apiKey, workspaceName1);

            UUID spanId = createSpan(factory.manufacturePojo(Span.class).toBuilder()
                    .projectName(projectName)
                    .build(), API_KEY, TEST_WORKSPACE);

            var item = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .spanId(spanId)
                    .traceId(traceId)
                    .source(DatasetItemSource.SPAN)
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(datasetId)
                    .build();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName1)
                    .put(Entity.json(batch))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(409);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(
                        "span workspace and dataset item workspace does not match");
            }
        }
    }

    private UUID createTrace(Trace trace, String apiKey, String workspaceName) {
        return traceResourceClient.createTrace(trace, apiKey, workspaceName);
    }

    private UUID createSpan(Span span, String apiKey, String workspaceName) {
        return spanResourceClient.createSpan(span, apiKey, workspaceName);
    }

    @Nested
    @DisplayName("Get dataset items {id}:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class GetDatasetItem {

        @Test
        @DisplayName("Success")
        void getDatasetItemById() {

            var item = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            getItemAndAssert(item, TEST_WORKSPACE, API_KEY);
        }

        @Test
        @DisplayName("when dataset item not found, then return 404")
        void getDatasetItemById__whenDatasetItemNotFound__thenReturn404() {
            String id = UUID.randomUUID().toString();

            var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path(id)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get();

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
            assertThat(actualResponse.hasEntity()).isTrue();
            assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset item not found");
        }

    }

    @Nested
    @DisplayName("Stream dataset items by {datasetId}:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class StreamDatasetItems {

        @Test
        @DisplayName("when streaming dataset items, then return items sorted by created date")
        void streamDataItems__whenStreamingDatasetItems__thenReturnItemsSortedByCreatedDate() {

            var items = IntStream.range(0, 10)
                    .mapToObj(i -> factory.manufacturePojo(DatasetItem.class))
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var streamRequest = DatasetItemStreamRequest.builder()
                    .datasetName(batch.datasetName())
                    .build();

            try (Response response = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("stream")
                    .request()
                    .accept(MediaType.APPLICATION_OCTET_STREAM)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(streamRequest))) {

                assertThat(response.getStatus()).isEqualTo(200);

                List<DatasetItem> actualItems = getStreamedItems(response);

                assertPage(items.reversed(), actualItems);
            }
        }

        @Test
        @DisplayName("when streaming dataset items with filters, then return items sorted by created date")
        void streamDataItems__whenStreamingDatasetItemsWithFilters__thenReturnItemsSortedByCreatedDate() {

            var items = IntStream.range(0, 5)
                    .mapToObj(i -> factory.manufacturePojo(DatasetItem.class))
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var streamRequest = DatasetItemStreamRequest.builder()
                    .datasetName(batch.datasetName())
                    .lastRetrievedId(items.reversed().get(1).id())
                    .build();

            try (Response response = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("stream")
                    .request()
                    .accept(MediaType.APPLICATION_OCTET_STREAM)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(streamRequest))) {

                assertThat(response.getStatus()).isEqualTo(200);

                List<DatasetItem> actualItems = getStreamedItems(response);

                assertPage(items.reversed().subList(2, 5), actualItems);
            }
        }

        @Test
        @DisplayName("when streaming has max steamLimit, then return items sorted by created date")
        void streamDataItems__whenStreamingHasMaxSize__thenReturnItemsSortedByCreatedDate() {

            var items = IntStream.range(0, 1000)
                    .mapToObj(i -> factory.manufacturePojo(DatasetItem.class).toBuilder()
                            .experimentItems(null)
                            .createdAt(null)
                            .lastUpdatedAt(null)
                            .build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            List<DatasetItem> expectedFirstPage = items.reversed().subList(0, 500);

            var streamRequest = DatasetItemStreamRequest.builder()
                    .datasetName(batch.datasetName()).build();

            try (Response response = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("stream")
                    .request()
                    .accept(MediaType.APPLICATION_OCTET_STREAM)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(streamRequest))) {

                assertThat(response.getStatus()).isEqualTo(200);

                List<DatasetItem> actualItems = getStreamedItems(response);

                assertPage(expectedFirstPage, actualItems);
            }

            streamRequest = DatasetItemStreamRequest.builder()
                    .datasetName(batch.datasetName())
                    .lastRetrievedId(expectedFirstPage.get(499).id())
                    .build();

            try (Response response = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("stream")
                    .request()
                    .accept(MediaType.APPLICATION_OCTET_STREAM)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(streamRequest))) {

                assertThat(response.getStatus()).isEqualTo(200);

                List<DatasetItem> actualItems = getStreamedItems(response);

                assertPage(items.reversed().subList(500, 1000), actualItems);
            }
        }
    }

    private DatasetItem getItemAndAssert(DatasetItem expectedDatasetItem, String workspaceName, String apiKey) {
        var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                .path("items")
                .path(expectedDatasetItem.id().toString())
                .request()
                .header(HttpHeaders.AUTHORIZATION, apiKey)
                .header(WORKSPACE_HEADER, workspaceName)
                .get();

        var actualEntity = actualResponse.readEntity(DatasetItem.class);
        assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);

        assertThat(actualEntity.id()).isEqualTo(expectedDatasetItem.id());
        assertThat(actualEntity).usingRecursiveComparison()
                .ignoringFields(IGNORED_FIELDS_DATA_ITEM)
                .isEqualTo(expectedDatasetItem);

        assertThat(actualEntity.createdAt()).isInThePast();
        assertThat(actualEntity.lastUpdatedAt()).isInThePast();

        return actualEntity;
    }

    @Nested
    @DisplayName("Patch dataset item:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class PatchDatasetItem {

        private Stream<Arguments> patchFieldScenarios() {
            return Stream.of(
                    // Test patching data field only
                    arguments("data field", (java.util.function.Function<DatasetItem, DatasetItem>) item -> {
                        var newData = Map.<String, JsonNode>of(
                                "updated_field", factory.manufacturePojo(JsonNode.class));
                        return DatasetItem.builder()
                                .data(newData)
                                .build();
                    }),
                    // Test patching source, traceId and spanId for SPAN source
                    arguments("source, traceId and spanId to SPAN",
                            (java.util.function.Function<DatasetItem, DatasetItem>) item -> DatasetItem
                                    .builder()
                                    .source(DatasetItemSource.SPAN)
                                    .traceId(GENERATOR.generate())
                                    .spanId(GENERATOR.generate())
                                    .build()),
                    // Test patching traceId only (keeping existing SPAN source and spanId)
                    arguments("traceId only",
                            (java.util.function.Function<DatasetItem, DatasetItem>) item -> DatasetItem
                                    .builder()
                                    .traceId(GENERATOR.generate())
                                    .build()),
                    // Test patching spanId only (keeping existing SPAN source and traceId)
                    arguments("spanId only", (java.util.function.Function<DatasetItem, DatasetItem>) item -> DatasetItem
                            .builder()
                            .spanId(GENERATOR.generate())
                            .build()));
        }

        @ParameterizedTest
        @MethodSource("patchFieldScenarios")
        @DisplayName("Success: patch different fields")
        void patchDatasetItem(String scenarioName, java.util.function.Function<DatasetItem, DatasetItem> patchBuilder) {
            // Create initial item with SPAN source
            var originalItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .source(DatasetItemSource.SPAN)
                    .traceId(GENERATOR.generate())
                    .spanId(GENERATOR.generate())
                    .tags(Set.of())
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(originalItem))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify original item exists
            var retrievedOriginal = getItemAndAssert(originalItem, TEST_WORKSPACE, API_KEY);

            // Build patch with only the fields we want to update
            var patchItem = patchBuilder.apply(originalItem);

            // Apply patch
            datasetResourceClient.patchDatasetItem(originalItem.id(), patchItem, API_KEY, TEST_WORKSPACE);

            // Retrieve patched item
            var patchedItem = datasetResourceClient.getDatasetItem(originalItem.id(), API_KEY, TEST_WORKSPACE);

            // Verify patched fields were updated
            if (patchItem.data() != null) {
                assertThat(patchedItem.data()).isEqualTo(patchItem.data());
            } else {
                assertThat(patchedItem.data()).isEqualTo(retrievedOriginal.data());
            }

            if (patchItem.source() != null) {
                assertThat(patchedItem.source()).isEqualTo(patchItem.source());
            } else {
                assertThat(patchedItem.source()).isEqualTo(retrievedOriginal.source());
            }

            if (patchItem.traceId() != null) {
                assertThat(patchedItem.traceId()).isEqualTo(patchItem.traceId());
            } else {
                assertThat(patchedItem.traceId()).isEqualTo(retrievedOriginal.traceId());
            }

            if (patchItem.spanId() != null) {
                assertThat(patchedItem.spanId()).isEqualTo(patchItem.spanId());
            } else {
                assertThat(patchedItem.spanId()).isEqualTo(retrievedOriginal.spanId());
            }

            // Verify read-only fields remain unchanged
            assertThat(patchedItem.id()).isEqualTo(originalItem.id());
            assertThat(patchedItem.datasetId()).isEqualTo(retrievedOriginal.datasetId());
            assertThat(patchedItem.createdBy()).isEqualTo(retrievedOriginal.createdBy());
            // Note: createdAt and lastUpdatedAt will be updated by ClickHouse on INSERT
        }

        @Test
        @DisplayName("when dataset item not found, then return 404")
        void patchDatasetItem__whenDatasetItemNotFound__thenReturn404() {
            var itemId = GENERATOR.generate();
            var patchItem = DatasetItem.builder()
                    .data(Map.of("field", factory.manufacturePojo(JsonNode.class)))
                    .build();

            try (var actualResponse = datasetResourceClient.callPatchDatasetItem(itemId, patchItem, API_KEY,
                    TEST_WORKSPACE)) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset item not found");
            }
        }

        @Test
        @DisplayName("Success: add tags to dataset item")
        void patchDatasetItem__whenAddingTags__thenSucceed() {
            // Create initial item without tags
            var originalItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of())
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(originalItem))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify original item has no tags
            var retrievedOriginal = datasetResourceClient.getDatasetItem(originalItem.id(), API_KEY, TEST_WORKSPACE);
            assertThat(retrievedOriginal.tags()).isEmpty();

            // Add tags via PATCH
            var tagsToAdd = Set.of("tag1", "tag2", "tag3");
            var patchItem = DatasetItem.builder()
                    .tags(tagsToAdd)
                    .build();

            datasetResourceClient.patchDatasetItem(originalItem.id(), patchItem, API_KEY, TEST_WORKSPACE);

            // Verify tags were added
            var patchedItem = datasetResourceClient.getDatasetItem(originalItem.id(), API_KEY, TEST_WORKSPACE);
            assertThat(patchedItem.tags()).containsExactlyInAnyOrderElementsOf(tagsToAdd);
        }

        @Test
        @DisplayName("Success: update tags on dataset item")
        void patchDatasetItem__whenUpdatingTags__thenSucceed() {
            // Create initial item with tags
            var initialTags = Set.of("tag1", "tag2");
            var originalItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(initialTags)
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(originalItem))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify original item has initial tags
            var retrievedOriginal = getItemAndAssert(originalItem, TEST_WORKSPACE, API_KEY);
            assertThat(retrievedOriginal.tags()).containsExactlyInAnyOrderElementsOf(initialTags);

            // Update tags via PATCH (add new tags)
            var updatedTags = Set.of("tag1", "tag2", "tag3", "tag4");
            var patchItem = DatasetItem.builder()
                    .tags(updatedTags)
                    .build();

            datasetResourceClient.patchDatasetItem(originalItem.id(), patchItem, API_KEY, TEST_WORKSPACE);

            // Verify tags were updated
            var patchedItem = datasetResourceClient.getDatasetItem(originalItem.id(), API_KEY, TEST_WORKSPACE);
            assertThat(patchedItem.tags()).containsExactlyInAnyOrderElementsOf(updatedTags);
        }

        @Test
        @DisplayName("Success: remove tags from dataset item")
        void patchDatasetItem__whenRemovingTags__thenSucceed() {
            // Create initial item with tags
            var initialTags = Set.of("tag1", "tag2", "tag3");
            var originalItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(initialTags)
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(originalItem))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify original item has initial tags
            var retrievedOriginal = getItemAndAssert(originalItem, TEST_WORKSPACE, API_KEY);
            assertThat(retrievedOriginal.tags()).containsExactlyInAnyOrderElementsOf(initialTags);

            // Remove some tags via PATCH
            var remainingTags = Set.of("tag1");
            var patchItem = DatasetItem.builder()
                    .tags(remainingTags)
                    .build();

            datasetResourceClient.patchDatasetItem(originalItem.id(), patchItem, API_KEY, TEST_WORKSPACE);

            // Verify tags were removed
            var patchedItem = datasetResourceClient.getDatasetItem(originalItem.id(), API_KEY, TEST_WORKSPACE);
            assertThat(patchedItem.tags()).containsExactlyInAnyOrderElementsOf(remainingTags);
        }

        @Test
        @DisplayName("Success: clear all tags from dataset item")
        void patchDatasetItem__whenClearingAllTags__thenSucceed() {
            // Create initial item with tags
            var initialTags = Set.of("tag1", "tag2", "tag3");
            var originalItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(initialTags)
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(originalItem))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify original item has initial tags
            var retrievedOriginal = getItemAndAssert(originalItem, TEST_WORKSPACE, API_KEY);
            assertThat(retrievedOriginal.tags()).containsExactlyInAnyOrderElementsOf(initialTags);

            // Clear all tags via PATCH
            var patchItem = DatasetItem.builder()
                    .tags(Set.of())
                    .build();

            datasetResourceClient.patchDatasetItem(originalItem.id(), patchItem, API_KEY, TEST_WORKSPACE);

            // Verify all tags were cleared
            var patchedItem = datasetResourceClient.getDatasetItem(originalItem.id(), API_KEY, TEST_WORKSPACE);
            assertThat(patchedItem.tags()).isEmpty();
        }
    }

    @Nested
    @DisplayName("Batch update dataset items:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class BatchUpdateDatasetItems {

        @Test
        @DisplayName("Success: batch add tags with merge")
        void batchUpdateDatasetItems__whenAddingTagsWithMerge__thenSucceed() {
            // Create items with different initial tags
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("existing1"))
                    .build();
            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("existing2"))
                    .build();
            var item3 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of())
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2, item3))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify initial state
            var retrieved1 = datasetResourceClient.getDatasetItem(item1.id(), API_KEY, TEST_WORKSPACE);
            var retrieved2 = datasetResourceClient.getDatasetItem(item2.id(), API_KEY, TEST_WORKSPACE);
            var retrieved3 = datasetResourceClient.getDatasetItem(item3.id(), API_KEY, TEST_WORKSPACE);
            assertThat(retrieved1.tags()).containsExactlyInAnyOrder("existing1");
            assertThat(retrieved2.tags()).containsExactlyInAnyOrder("existing2");
            assertThat(retrieved3.tags()).isEmpty();

            // Batch add tags with merge
            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .ids(Set.of(item1.id(), item2.id(), item3.id()))
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("newtag"))
                            .build())
                    .mergeTags(true)
                    .build();

            datasetResourceClient.batchUpdateDatasetItems(batchUpdate, API_KEY, TEST_WORKSPACE);

            // Verify tags were merged
            var updated1 = datasetResourceClient.getDatasetItem(item1.id(), API_KEY, TEST_WORKSPACE);
            var updated2 = datasetResourceClient.getDatasetItem(item2.id(), API_KEY, TEST_WORKSPACE);
            var updated3 = datasetResourceClient.getDatasetItem(item3.id(), API_KEY, TEST_WORKSPACE);
            assertThat(updated1.tags()).containsExactlyInAnyOrder("existing1", "newtag");
            assertThat(updated2.tags()).containsExactlyInAnyOrder("existing2", "newtag");
            assertThat(updated3.tags()).containsExactlyInAnyOrder("newtag");
        }

        @Test
        @DisplayName("Success: batch replace tags without merge")
        void batchUpdateDatasetItems__whenReplacingTagsWithoutMerge__thenSucceed() {
            // Create items with initial tags
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("old1", "old2"))
                    .build();
            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("old3", "old4"))
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify initial state
            var retrieved1 = datasetResourceClient.getDatasetItem(item1.id(), API_KEY, TEST_WORKSPACE);
            var retrieved2 = datasetResourceClient.getDatasetItem(item2.id(), API_KEY, TEST_WORKSPACE);
            assertThat(retrieved1.tags()).containsExactlyInAnyOrder("old1", "old2");
            assertThat(retrieved2.tags()).containsExactlyInAnyOrder("old3", "old4");

            // Batch replace tags without merge
            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .ids(Set.of(item1.id(), item2.id()))
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("new1", "new2"))
                            .build())
                    .mergeTags(false)
                    .build();

            datasetResourceClient.batchUpdateDatasetItems(batchUpdate, API_KEY, TEST_WORKSPACE);

            // Verify tags were replaced
            var updated1 = datasetResourceClient.getDatasetItem(item1.id(), API_KEY, TEST_WORKSPACE);
            var updated2 = datasetResourceClient.getDatasetItem(item2.id(), API_KEY, TEST_WORKSPACE);
            assertThat(updated1.tags()).containsExactlyInAnyOrder("new1", "new2");
            assertThat(updated2.tags()).containsExactlyInAnyOrder("new1", "new2");
        }

        @Test
        @DisplayName("Error: batch update with empty ids")
        void batchUpdateDatasetItems__whenEmptyIds__thenBadRequest() {
            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .ids(Set.of())
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("tag"))
                            .build())
                    .mergeTags(true)
                    .build();

            try (var actualResponse = datasetResourceClient.callBatchUpdateDatasetItems(batchUpdate, API_KEY,
                    TEST_WORKSPACE)) {
                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
            }
        }

        @Test
        @DisplayName("Error: batch update exceeds max size")
        void batchUpdateDatasetItems__whenExceedsMaxSize__thenBadRequest() {
            // Create more than 1000 IDs
            var tooManyIds = new HashSet<UUID>();
            for (int i = 0; i < 1001; i++) {
                tooManyIds.add(UUID.randomUUID());
            }

            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .ids(tooManyIds)
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("tag"))
                            .build())
                    .mergeTags(true)
                    .build();

            try (var actualResponse = datasetResourceClient.callBatchUpdateDatasetItems(batchUpdate, API_KEY,
                    TEST_WORKSPACE)) {
                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
            }
        }

        @Test
        @DisplayName("Success: batch update by filters with merge tags")
        void batchUpdateDatasetItems__whenUsingFiltersWithMerge__thenSucceed() {
            // Create items with different tags
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("include", "tag1"))
                    .build();
            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("include", "tag2"))
                    .build();
            var item3 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("exclude", "tag3"))
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2, item3))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify initial state
            var retrieved1 = datasetResourceClient.getDatasetItem(item1.id(), API_KEY, TEST_WORKSPACE);
            var retrieved2 = datasetResourceClient.getDatasetItem(item2.id(), API_KEY, TEST_WORKSPACE);
            var retrieved3 = datasetResourceClient.getDatasetItem(item3.id(), API_KEY, TEST_WORKSPACE);
            assertThat(retrieved1.tags()).containsExactlyInAnyOrder("include", "tag1");
            assertThat(retrieved2.tags()).containsExactlyInAnyOrder("include", "tag2");
            assertThat(retrieved3.tags()).containsExactlyInAnyOrder("exclude", "tag3");

            // Create filter to match items with "include" tag
            var filter = new DatasetItemFilter(DatasetItemField.TAGS, Operator.CONTAINS, null, "include");

            // Batch update by filters with merge
            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .filters(List.of(filter))
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("newtag"))
                            .build())
                    .mergeTags(true)
                    .build();

            datasetResourceClient.batchUpdateDatasetItems(batchUpdate, API_KEY, TEST_WORKSPACE);

            // Verify only items matching filters were updated
            var updated1 = datasetResourceClient.getDatasetItem(item1.id(), API_KEY, TEST_WORKSPACE);
            var updated2 = datasetResourceClient.getDatasetItem(item2.id(), API_KEY, TEST_WORKSPACE);
            var updated3 = datasetResourceClient.getDatasetItem(item3.id(), API_KEY, TEST_WORKSPACE);
            assertThat(updated1.tags()).containsExactlyInAnyOrder("include", "tag1", "newtag");
            assertThat(updated2.tags()).containsExactlyInAnyOrder("include", "tag2", "newtag");
            assertThat(updated3.tags()).containsExactlyInAnyOrder("exclude", "tag3"); // unchanged
        }

        @Test
        @DisplayName("Success: batch update by filters automatically merges tags even when mergeTags is false")
        void batchUpdateDatasetItems__whenUsingFiltersWithMergeFalse__thenAutoMerges() {
            // Create items with different tags
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("include", "tag1"))
                    .build();
            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("include", "tag2"))
                    .build();
            var item3 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("other"))
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2, item3))
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Verify initial state
            var retrieved1 = datasetResourceClient.getDatasetItem(item1.id(), API_KEY, TEST_WORKSPACE);
            var retrieved2 = datasetResourceClient.getDatasetItem(item2.id(), API_KEY, TEST_WORKSPACE);
            var retrieved3 = datasetResourceClient.getDatasetItem(item3.id(), API_KEY, TEST_WORKSPACE);
            assertThat(retrieved1.tags()).containsExactlyInAnyOrder("include", "tag1");
            assertThat(retrieved2.tags()).containsExactlyInAnyOrder("include", "tag2");
            assertThat(retrieved3.tags()).containsExactlyInAnyOrder("other");

            // Create filter to match items with "include" tag
            var filter = new DatasetItemFilter(DatasetItemField.TAGS, Operator.CONTAINS, null, "include");

            // Batch update by filters with mergeTags=false (should auto-set to true)
            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .filters(List.of(filter))
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("newtag"))
                            .build())
                    .mergeTags(false) // This will be automatically set to true
                    .build();

            datasetResourceClient.batchUpdateDatasetItems(batchUpdate, API_KEY, TEST_WORKSPACE);

            // Verify tags were merged (not replaced) for items matching filters
            var updated1 = datasetResourceClient.getDatasetItem(item1.id(), API_KEY, TEST_WORKSPACE);
            var updated2 = datasetResourceClient.getDatasetItem(item2.id(), API_KEY, TEST_WORKSPACE);
            var updated3 = datasetResourceClient.getDatasetItem(item3.id(), API_KEY, TEST_WORKSPACE);
            assertThat(updated1.tags()).containsExactlyInAnyOrder("include", "tag1", "newtag");
            assertThat(updated2.tags()).containsExactlyInAnyOrder("include", "tag2", "newtag");
            assertThat(updated3.tags()).containsExactlyInAnyOrder("other"); // Unchanged
        }

        @Test
        @DisplayName("Error: batch update with both ids and filters")
        void batchUpdateDatasetItems__whenBothIdsAndFilters__thenBadRequest() {
            var filter = new DatasetItemFilter(DatasetItemField.TAGS, Operator.CONTAINS, null, "tag");

            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .ids(Set.of(UUID.randomUUID()))
                    .filters(List.of(filter))
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("tag"))
                            .build())
                    .mergeTags(true)
                    .build();

            try (var actualResponse = datasetResourceClient.callBatchUpdateDatasetItems(batchUpdate, API_KEY,
                    TEST_WORKSPACE)) {
                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
            }
        }

        @Test
        @DisplayName("Error: batch update with neither ids nor filters")
        void batchUpdateDatasetItems__whenNeitherIdsNorFilters__thenBadRequest() {
            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("tag"))
                            .build())
                    .mergeTags(true)
                    .build();

            try (var actualResponse = datasetResourceClient.callBatchUpdateDatasetItems(batchUpdate, API_KEY,
                    TEST_WORKSPACE)) {
                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
            }
        }

        @Test
        @DisplayName("Error: batch update with empty filters")
        void batchUpdateDatasetItems__whenEmptyFilters__thenBadRequest() {
            var batchUpdate = DatasetItemBatchUpdate.builder()
                    .filters(List.of())
                    .update(DatasetItemUpdate.builder()
                            .tags(Set.of("tag"))
                            .build())
                    .mergeTags(true)
                    .build();

            try (var actualResponse = datasetResourceClient.callBatchUpdateDatasetItems(batchUpdate, API_KEY,
                    TEST_WORKSPACE)) {
                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
            }
        }
    }

    @Nested
    @DisplayName("Delete items:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class DeleteDatasetItems {

        @Test
        @DisplayName("Success")
        void deleteDatasetItem() {
            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(null)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var itemIds = items.stream().map(DatasetItem::id).toList();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("delete")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(new DatasetItemsDelete(itemIds)))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
            }

            for (var item : items) {
                var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                        .path("items")
                        .path(item.id().toString())
                        .request()
                        .header(HttpHeaders.AUTHORIZATION, API_KEY)
                        .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                        .get();

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(404);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains("Dataset item not found");
            }
        }

        @Test
        @DisplayName("when dataset item does not exists, then return no content")
        void deleteDatasetItem__whenDatasetItemDoesNotExists__thenReturnNoContent() {
            var id = UUID.randomUUID().toString();
            var itemIds = List.of(UUID.fromString(id));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("delete")
                    .request()
                    .accept(MediaType.APPLICATION_JSON_TYPE)
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(new DatasetItemsDelete(itemIds)))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                assertThat(actualResponse.hasEntity()).isFalse();
            }
        }

        @ParameterizedTest
        @MethodSource("invalidDatasetItemBatches")
        @DisplayName("when dataset item batch is not valid, then return 422")
        void deleteDatasetItem__whenDatasetItemIsNotValid__thenReturn422(List<UUID> itemIds, String errorMessage) {
            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("delete")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .post(Entity.json(new DatasetItemsDelete(itemIds)))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(422);
                assertThat(actualResponse.hasEntity()).isTrue();
                assertThat(actualResponse.readEntity(ErrorMessage.class).errors()).contains(errorMessage);
            }
        }

        public Stream<Arguments> invalidDatasetItemBatches() {
            return Stream.of(
                    arguments(List.of(),
                            "itemIds size must be between 1 and 1000"),
                    arguments(null,
                            "itemIds must not be null"),
                    arguments(IntStream.range(1, 10001).mapToObj(__ -> UUID.randomUUID()).toList(),
                            "itemIds size must be between 1 and 1000"));
        }
    }

    @Nested
    @DisplayName("Get dataset items by dataset id:")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class GetDatasetItemsByDatasetId {

        @Test
        @DisplayName("Success")
        void getDatasetItemsByDatasetId() {

            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .build();

            List<Map<String, JsonNode>> data = batch.items()
                    .stream()
                    .map(DatasetItem::data)
                    .toList();

            Set<Column> columns = getColumns(data);

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var actualEntity = datasetResourceClient.getDatasetItems(datasetId, Map.of(), API_KEY, TEST_WORKSPACE);

            assertDatasetItemPage(actualEntity, items.reversed(), columns, 1);
        }

        @Test
        @DisplayName("when defining page size, then return page with limit respected")
        void getDatasetItemsByDatasetId__whenDefiningPageSize__thenReturnPageWithLimitRespected() {

            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .build();

            List<Map<String, JsonNode>> data = batch.items()
                    .stream()
                    .map(DatasetItem::data)
                    .toList();

            Set<Column> columns = getColumns(data);

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var actualEntity = datasetResourceClient.getDatasetItems(datasetId, Map.of("size", 1), API_KEY,
                    TEST_WORKSPACE);

            List<DatasetItem> expectedContent = List.of(items.reversed().getFirst());

            assertDatasetItemPage(actualEntity, expectedContent, 5, columns, 1);
        }

        @Test
        @DisplayName("when items were updated, then return correct items count")
        void getDatasetItemsByDatasetId__whenItemsWereUpdated__thenReturnCorrectItemsCount() {

            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class);

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var updatedItems = items
                    .stream()
                    .map(item -> item.toBuilder()
                            .data(Map.of(factory.manufacturePojo(String.class),
                                    factory.manufacturePojo(JsonNode.class)))
                            .build())
                    .toList();

            var updatedBatch = batch.toBuilder()
                    .items(updatedItems)
                    .build();

            putAndAssert(updatedBatch, TEST_WORKSPACE, API_KEY);

            List<Map<String, JsonNode>> data = updatedBatch.items()
                    .stream()
                    .map(DatasetItem::data)
                    .toList();

            Set<Column> columns = getColumns(data);

            var actualEntity = datasetResourceClient.getDatasetItems(datasetId, Map.of(), API_KEY, TEST_WORKSPACE);

            assertDatasetItemPage(actualEntity, updatedItems.reversed(), columns, 1);
        }

        @Test
        @DisplayName("when items have data with same keys and different types, then return columns types and count")
        void getDatasetItemsByDatasetId__whenItemsHaveDataWithSameKeysAndDifferentTypes__thenReturnColumnsTypesAndCount() {

            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            var item = factory.manufacturePojo(DatasetItem.class);

            var item2 = item.toBuilder()
                    .id(factory.manufacturePojo(UUID.class))
                    .data(item.data()
                            .keySet()
                            .stream()
                            .map(key -> Map.entry(key, NullNode.getInstance()))
                            .collect(toMap(Map.Entry::getKey, Map.Entry::getValue)))
                    .build();

            var item3 = item.toBuilder()
                    .id(factory.manufacturePojo(UUID.class))
                    .data(item.data()
                            .keySet()
                            .stream()
                            .map(key -> Map.entry(key, TextNode.valueOf(RandomStringUtils.randomAlphanumeric(10))))
                            .collect(toMap(Map.Entry::getKey, Map.Entry::getValue)))
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item, item2, item3))
                    .datasetId(datasetId)
                    .build();

            List<Map<String, JsonNode>> data = batch.items()
                    .stream()
                    .map(DatasetItem::data)
                    .toList();

            Set<Column> columns = getColumns(data);

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var actualEntity = datasetResourceClient.getDatasetItems(datasetId, Map.of(), API_KEY, TEST_WORKSPACE);

            assertDatasetItemPage(actualEntity, batch.items().reversed(), columns, 1);

        }

        @ParameterizedTest
        @MethodSource("com.comet.opik.api.resources.utils.ImageTruncationArgProvider#provideTestArguments")
        void getDatasetItemsByDatasetId_withTruncation(JsonNode original, JsonNode expected, boolean truncate) {

            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            var items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class).stream()
                    .map(item -> item.toBuilder().data(ImmutableMap.of("image", original)).build())
                    .toList();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .build();

            List<Map<String, JsonNode>> data = batch.items()
                    .stream()
                    .map(DatasetItem::data)
                    .toList();

            Set<Column> columns = getColumns(data);

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var expectedDatasetItems = items.stream()
                    .map(item -> item.toBuilder().data(ImmutableMap.of("image", expected)).build())
                    .toList().reversed();

            var actualEntity = datasetResourceClient.getDatasetItems(datasetId, Map.of("truncate", truncate), API_KEY,
                    TEST_WORKSPACE);

            assertDatasetItemPage(actualEntity, expectedDatasetItems, columns, 1);
        }

        Stream<Arguments> dataFieldFilterOperators() {
            return Stream.of(
                    // CONTAINS - matches when value contains the search term
                    Arguments.of(
                            Named.of("CONTAINS operator", Operator.CONTAINS),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode("search for " + searchKey + " model"),
                                    "type", new TextNode("question")),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode("completely different content"),
                                    "type", new TextNode("recipe"))),

                    // NOT_CONTAINS - matches when value does NOT contain the search term
                    Arguments.of(
                            Named.of("NOT_CONTAINS operator", Operator.NOT_CONTAINS),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode("completely different content"),
                                    "type", new TextNode("recipe")),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode("search for " + searchKey + " model"),
                                    "type", new TextNode("question"))),

                    // STARTS_WITH - matches when value starts with the search term
                    Arguments.of(
                            Named.of("STARTS_WITH operator", Operator.STARTS_WITH),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode(searchKey + " is the beginning"),
                                    "type", new TextNode("question")),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode("this does not start with " + searchKey),
                                    "type", new TextNode("recipe"))),

                    // ENDS_WITH - matches when value ends with the search term
                    Arguments.of(
                            Named.of("ENDS_WITH operator", Operator.ENDS_WITH),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode("this ends with " + searchKey),
                                    "type", new TextNode("question")),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode(searchKey + " is at the start not end"),
                                    "type", new TextNode("recipe"))),

                    // EQUAL - matches when value equals the search term (case insensitive)
                    Arguments.of(
                            Named.of("EQUAL operator", Operator.EQUAL),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode(searchKey),
                                    "type", new TextNode("question")),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode(searchKey + " extra text"),
                                    "type", new TextNode("recipe"))),

                    // NOT_EQUAL - matches when value does NOT equal the search term
                    Arguments.of(
                            Named.of("NOT_EQUAL operator", Operator.NOT_EQUAL),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode(searchKey + " extra text"),
                                    "type", new TextNode("question")),
                            (Function<String, Map<String, JsonNode>>) searchKey -> Map.of(
                                    "query", new TextNode(searchKey),
                                    "type", new TextNode("recipe"))));
        }

        @ParameterizedTest
        @MethodSource("dataFieldFilterOperators")
        @DisplayName("when filtering data field with operator, then return matching items")
        void getDatasetItemsByDatasetId__whenFilteringDataFieldWithOperator__thenReturnMatchingItems(
                Operator operator,
                Function<String, Map<String, JsonNode>> matchingDataSupplier,
                Function<String, Map<String, JsonNode>> nonMatchingDataSupplier) {

            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            var searchKey = RandomStringUtils.secure().nextAlphabetic(8);
            var matchingItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .data(matchingDataSupplier.apply(searchKey))
                    .build();

            var nonMatchingItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .data(nonMatchingDataSupplier.apply(searchKey))
                    .build();

            var items = List.of(matchingItem, nonMatchingItem);
            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(items)
                    .datasetId(datasetId)
                    .build();

            List<Map<String, JsonNode>> data = batch.items()
                    .stream()
                    .map(DatasetItem::data)
                    .toList();

            Set<Column> columns = getColumns(data);

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            var filter = new DatasetItemFilter(DatasetItemField.DATA, operator, "query", searchKey);

            var actualEntity = datasetResourceClient.getDatasetItems(datasetId,
                    Map.of("filters", toURLEncodedQueryParam(List.of(filter))), API_KEY, TEST_WORKSPACE);

            assertDatasetItemPage(actualEntity, List.of(matchingItem), columns, 1);
        }

        // Helper method to create dataset items with content
        private DatasetItem createDatasetItem(String content) {
            return factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .data(Map.of("content", new TextNode(content)))
                    .build();
        }

        // Record to hold test scenario data
        private record SearchTestScenario(
                List<DatasetItem> items,
                String searchTerm,
                int page,
                int size,
                int expectedTotalCount,
                List<DatasetItem> expectedItems) {
        }

        @Test
        @DisplayName("Success: filter dataset items by tags")
        void getDatasetItemsByDatasetId__whenFilteringByTags__thenSucceed() {
            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            // Create items with different tags
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag1", "tag2"))
                    .build();
            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag2", "tag3"))
                    .build();
            var item3 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag3", "tag4"))
                    .build();
            var item4 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag5"))
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2, item3, item4))
                    .datasetId(datasetId)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Filter by tag2 - should return item1 and item2
            var filter = new DatasetItemFilter(DatasetItemField.TAGS, Operator.CONTAINS, null, "tag2");
            var actualEntity = datasetResourceClient.getDatasetItems(datasetId,
                    Map.of("filters", toURLEncodedQueryParam(List.of(filter))), API_KEY, TEST_WORKSPACE);

            List<DatasetItem> expectedContent = List.of(item2, item1); // reversed order

            assertThat(actualEntity.content()).hasSize(2);
            assertThat(actualEntity.total()).isEqualTo(2);
            assertThat(actualEntity.content().getFirst().id()).isEqualTo(expectedContent.getFirst().id());
            assertThat(actualEntity.content().get(1).id()).isEqualTo(expectedContent.get(1).id());
        }

        @Test
        @DisplayName("Success: filter dataset items by multiple tags")
        void getDatasetItemsByDatasetId__whenFilteringByMultipleTags__thenSucceed() {
            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            // Create items with different tags
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag1", "tag2"))
                    .build();
            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag2", "tag3"))
                    .build();
            var item3 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag3", "tag4"))
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2, item3))
                    .datasetId(datasetId)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Filter by tag3 - should return item2 and item3
            var filter = new DatasetItemFilter(DatasetItemField.TAGS, Operator.CONTAINS, null, "tag3");
            var actualEntity = datasetResourceClient.getDatasetItems(datasetId,
                    Map.of("filters", toURLEncodedQueryParam(List.of(filter))), API_KEY, TEST_WORKSPACE);

            assertThat(actualEntity.content()).hasSize(2);
            assertThat(actualEntity.total()).isEqualTo(2);
        }

        @Test
        @DisplayName("Success: filter dataset items by non-existent tag")
        void getDatasetItemsByDatasetId__whenFilteringByNonExistentTag__thenReturnEmpty() {
            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            // Create items with different tags
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag1", "tag2"))
                    .build();
            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag3", "tag4"))
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2))
                    .datasetId(datasetId)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Filter by non-existent tag
            var filter = new DatasetItemFilter(DatasetItemField.TAGS, Operator.CONTAINS, null, "nonexistent");
            var actualEntity = datasetResourceClient.getDatasetItems(datasetId,
                    Map.of("filters", toURLEncodedQueryParam(List.of(filter))), API_KEY, TEST_WORKSPACE);

            assertThat(actualEntity.content()).isEmpty();
            assertThat(actualEntity.total()).isEqualTo(0);
        }

        @Test
        @DisplayName("Success: filter dataset items without tags")
        void getDatasetItemsByDatasetId__whenFilteringItemsWithoutTags__thenSucceed() {
            UUID datasetId = createAndAssert(factory.manufacturePojo(Dataset.class).toBuilder()
                    .id(null)
                    .build());

            // Create items, some with tags and some without
            var item1 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of("tag1"))
                    .build();
            var item2 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(null)
                    .build();
            var item3 = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .tags(Set.of())
                    .build();

            var batch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .items(List.of(item1, item2, item3))
                    .datasetId(datasetId)
                    .build();

            putAndAssert(batch, TEST_WORKSPACE, API_KEY);

            // Get all items without filtering - should return all 3
            var allItems = datasetResourceClient.getDatasetItems(datasetId, Map.of(), API_KEY, TEST_WORKSPACE);
            assertThat(allItems.content()).hasSize(3);

            // Filter by tag1 - should return only item1
            var filter = new DatasetItemFilter(DatasetItemField.TAGS, Operator.CONTAINS, null, "tag1");
            var filteredItems = datasetResourceClient.getDatasetItems(datasetId,
                    Map.of("filters", toURLEncodedQueryParam(List.of(filter))), API_KEY, TEST_WORKSPACE);

            assertThat(filteredItems.content()).hasSize(1);
            assertThat(filteredItems.content().getFirst().id()).isEqualTo(item1.id());
        }
    }

    private void assertDatasetItemPage(DatasetItemPage actualPage, List<DatasetItem> expected, Set<Column> columns,
            int page) {
        assertDatasetItemPage(actualPage, expected, expected.size(), columns, page);
    }

    private void assertDatasetItemPage(DatasetItemPage actualPage, List<DatasetItem> expected, int total,
            Set<Column> columns, int page) {
        assertThat(actualPage.size()).isEqualTo(expected.size());
        assertThat(actualPage.content()).hasSize(expected.size());
        assertThat(actualPage.page()).isEqualTo(page);
        assertThat(actualPage.total()).isEqualTo(total);
        assertThat(actualPage.columns()).isEqualTo(columns);

        assertPage(expected, actualPage.content());
    }

    private void assertPage(List<DatasetItem> expectedItems, List<DatasetItem> actualItems) {

        List<String> ignoredFields = new ArrayList<>(Arrays.asList(IGNORED_FIELDS_DATA_ITEM));
        ignoredFields.add("data");

        assertThat(actualItems)
                .usingRecursiveFieldByFieldElementComparatorIgnoringFields(ignoredFields.toArray(String[]::new))
                .isEqualTo(expectedItems);

        assertThat(actualItems).hasSize(expectedItems.size());
        for (int i = 0; i < actualItems.size(); i++) {
            var actualDatasetItem = actualItems.get(i);
            var expectedDatasetItem = expectedItems.get(i);

            assertThat(actualDatasetItem.data()).isEqualTo(expectedDatasetItem.data());
        }
    }

    @Nested
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class FindDatasetItemsWithExperimentItems {

        @Test
        void find() {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            // Creating two traces with input, output and scores
            var trace1 = factory.manufacturePojo(Trace.class);
            createAndAssert(trace1, workspaceName, apiKey);

            var trace2 = factory.manufacturePojo(Trace.class);
            createAndAssert(trace2, workspaceName, apiKey);
            var traces = List.of(trace1, trace2);

            // Creating 5 scores peach each of the two traces above
            var scores1 = PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItem.class)
                    .stream()
                    .map(feedbackScoreBatchItem -> feedbackScoreBatchItem.toBuilder()
                            .id(trace1.id())
                            .projectName(trace1.projectName())
                            .value(factory.manufacturePojo(BigDecimal.class))
                            .build())
                    .toList();

            var scores2 = PodamFactoryUtils.manufacturePojoList(factory, FeedbackScoreBatchItem.class)
                    .stream()
                    .map(feedbackScoreBatchItem -> feedbackScoreBatchItem.toBuilder()
                            .id(trace2.id())
                            .projectName(trace2.projectName())
                            .value(factory.manufacturePojo(BigDecimal.class))
                            .build())
                    .toList();

            var traceIdToScoresMap = Stream.concat(scores1.stream(), scores2.stream())
                    .collect(groupingBy(FeedbackScoreItem::id));

            // When storing the scores in batch, adding some more unrelated random ones
            var feedbackScoreBatch = factory.manufacturePojo(FeedbackScoreBatch.class);
            feedbackScoreBatch = feedbackScoreBatch.toBuilder()
                    .scores(Stream.concat(
                            feedbackScoreBatch.scores().stream(),
                            traceIdToScoresMap.values().stream().flatMap(List::stream)).toList())
                    .build();

            createScoreAndAssert(feedbackScoreBatch, apiKey, workspaceName);

            // Add comments to random trace
            List<Comment> expectedComments = IntStream.range(0, 5)
                    .mapToObj(i -> traceResourceClient.generateAndCreateComment(trace1.id(), apiKey, workspaceName,
                            201))
                    .toList();

            // Creating a trace without input, output and scores
            var traceMissingFields = factory.manufacturePojo(Trace.class).toBuilder()
                    .input(null)
                    .output(null)
                    .build();
            createAndAssert(traceMissingFields, workspaceName, apiKey);

            // Creating the dataset
            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            // Creating 5 dataset items for the dataset above
            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            // Creating 5 different experiment ids
            var expectedDatasetItems = datasetItemBatch.items().subList(0, 4).reversed();
            var experimentIds = IntStream.range(0, 5).mapToObj(__ -> GENERATOR.generate()).toList();

            // Dataset items 0 and 1 cover the general case.
            // Per each dataset item there are 10 experiment items, so 2 experiment items per each of the 5 experiments.
            // The first 5 experiment items are related to trace 1, the other 5 to trace 2.
            var datasetItemIdToExperimentItemMap = expectedDatasetItems.subList(0, 2).stream()
                    .flatMap(datasetItem -> IntStream.range(0, 10)
                            .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                                    .experimentId(experimentIds.get(i / 2))
                                    .datasetItemId(datasetItem.id())
                                    .traceId(traces.get(i / 5).id())
                                    .input(traces.get(i / 5).input())
                                    .output(traces.get(i / 5).output())
                                    .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(
                                            traces.get(i / 5).startTime(), traces.get(i / 5).endTime()))
                                    .feedbackScores(traceIdToScoresMap.get(traces.get(i / 5).id()).stream()
                                            .map(FeedbackScoreMapper.INSTANCE::toFeedbackScore)
                                            .toList())
                                    .usage(null)
                                    .totalEstimatedCost(null)
                                    .build()))
                    .collect(groupingBy(ExperimentItem::datasetItemId));

            // Dataset item 2 covers the case of experiments items related to a trace without input, output and scores.
            // It also has 2 experiment items per each of the 5 experiments.
            datasetItemIdToExperimentItemMap.put(expectedDatasetItems.get(2).id(), experimentIds.stream()
                    .flatMap(experimentId -> IntStream.range(0, 2)
                            .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                                    .experimentId(experimentId)
                                    .datasetItemId(expectedDatasetItems.get(2).id())
                                    .traceId(traceMissingFields.id())
                                    .input(traceMissingFields.input())
                                    .output(traceMissingFields.output())
                                    .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(
                                            traceMissingFields.startTime(), traceMissingFields.endTime()))
                                    .feedbackScores(null)
                                    .usage(null)
                                    .totalEstimatedCost(null)
                                    .build()))
                    .toList());

            // Dataset item 3 covers the case of experiments items related to an un-existing trace id.
            // It also has 2 experiment items per each of the 5 experiments.
            datasetItemIdToExperimentItemMap.put(expectedDatasetItems.get(3).id(), experimentIds.stream()
                    .flatMap(experimentId -> IntStream.range(0, 2)
                            .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                                    .experimentId(experimentId)
                                    .datasetItemId(expectedDatasetItems.get(3).id())
                                    .input(null)
                                    .output(null)
                                    .feedbackScores(null)
                                    .usage(null)
                                    .totalEstimatedCost(null)
                                    .duration(null)
                                    .traceVisibilityMode(null)
                                    .build()))
                    .toList());

            // Dataset item 4 covers the case of not matching experiment items.

            // When storing the experiment items in batch, adding some more unrelated random ones
            var experimentItemsBatch = factory.manufacturePojo(ExperimentItemsBatch.class);
            experimentItemsBatch = experimentItemsBatch.toBuilder()
                    .experimentItems(Stream.concat(experimentItemsBatch.experimentItems().stream()
                            .map(item -> item.toBuilder()
                                    .usage(null)
                                    .totalEstimatedCost(null)
                                    .duration(null)
                                    .traceVisibilityMode(null)
                                    .build()),
                            datasetItemIdToExperimentItemMap.values().stream().flatMap(Collection::stream))
                            .collect(toUnmodifiableSet()))
                    .build();
            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            List<Map<String, JsonNode>> data = datasetItemBatch.items().stream()
                    .map(DatasetItem::data)
                    .toList();

            Set<Column> columns = getColumns(data);

            var page = 1;
            var pageSize = 5;
            // Filtering by experiments 1 and 3.
            var experimentIdsQueryParm = JsonUtils
                    .writeValueAsString(List.of(experimentIds.get(1), experimentIds.get(3)));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("page", page)
                    .queryParam("size", pageSize)
                    .queryParam("experiment_ids", experimentIdsQueryParm)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                var actualPage = actualResponse.readEntity(DatasetItemPage.class);

                assertThat(actualPage.page()).isEqualTo(page);
                assertThat(actualPage.size()).isEqualTo(expectedDatasetItems.size());
                assertThat(actualPage.total()).isEqualTo(expectedDatasetItems.size());
                assertThat(actualPage.columns()).isEqualTo(columns);

                var actualDatasetItems = actualPage.content();

                assertPage(expectedDatasetItems, actualPage.content());

                for (var i = 0; i < actualDatasetItems.size(); i++) {
                    var actualDatasetItem = actualDatasetItems.get(i);
                    var expectedDatasetItem = expectedDatasetItems.get(i);

                    // Filtering by those related to experiments 1 and 3
                    var experimentItems = datasetItemIdToExperimentItemMap.get(expectedDatasetItem.id());
                    var expectedExperimentItems = List.of(
                            experimentItems.get(2),
                            experimentItems.get(3),
                            experimentItems.get(6),
                            experimentItems.get(7)).reversed();

                    assertThat(actualDatasetItem.experimentItems())
                            .usingRecursiveComparison()
                            .withComparatorForType(StatsUtils::bigDecimalComparator, BigDecimal.class)
                            .withComparatorForFields(StatsUtils::closeToEpsilonComparator, "duration")
                            .ignoringFields(IGNORED_FIELDS_LIST)
                            .isEqualTo(expectedExperimentItems);

                    for (var j = 0; j < actualDatasetItem.experimentItems().size(); j++) {
                        var actualExperimentItem = assertFeedbackScoresIgnoredFieldsAndSetThemToNull(
                                actualDatasetItem.experimentItems().get(j), USER);
                        var expectedExperimentItem = expectedExperimentItems.get(j);

                        assertThat(actualExperimentItem.feedbackScores())
                                .usingRecursiveComparison()
                                .withComparatorForType(BigDecimal::compareTo, BigDecimal.class)
                                .ignoringCollectionOrder()
                                .isEqualTo(expectedExperimentItem.feedbackScores());

                        assertThat(actualExperimentItem.createdAt())
                                .isAfter(expectedExperimentItem.createdAt());
                        assertThat(actualExperimentItem.lastUpdatedAt())
                                .isAfter(expectedExperimentItem.lastUpdatedAt());

                        assertThat(actualExperimentItem.createdBy())
                                .isEqualTo(USER);
                        assertThat(actualExperimentItem.lastUpdatedBy())
                                .isEqualTo(USER);

                        // Check comments
                        if (actualExperimentItem.traceId().equals(trace1.id())) {
                            assertThat(expectedComments)
                                    .usingRecursiveComparison()
                                    .ignoringFields(IGNORED_FIELDS_COMMENTS)
                                    .isEqualTo(actualExperimentItem.comments());
                        }
                    }

                    assertThat(actualDatasetItem.createdAt()).isAfter(expectedDatasetItem.createdAt());
                    assertThat(actualDatasetItem.lastUpdatedAt()).isAfter(expectedDatasetItem.lastUpdatedAt());
                }
            }
        }

        @Test
        void find__whenExperimentsHaveSpansWithLLMCalls__thenIncludeSpanData() {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            // Creating two traces with input, output and scores

            var trace1 = factory.manufacturePojo(Trace.class);
            createAndAssert(trace1, workspaceName, apiKey);

            var trace2 = factory.manufacturePojo(Trace.class);
            createAndAssert(trace2, workspaceName, apiKey);

            var traces = List.of(trace1, trace2);

            Map<UUID, List<Span>> spansMap = traces.stream().map(trace -> {
                Span span1 = createSpan(trace, apiKey, workspaceName);
                Span span2 = createSpan(trace, apiKey, workspaceName);

                return Map.entry(trace.id(), List.of(span1, span2));
            }).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

            // Creating dataset and experiment items

            var dataset = factory.manufacturePojo(Dataset.class);

            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            // Creating 5 dataset items for the dataset above
            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetName(dataset.name())
                    .datasetId(datasetId)
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            var experiment = factory.manufacturePojo(Experiment.class).toBuilder()
                    .datasetName(dataset.name())
                    .promptVersion(null)
                    .promptVersions(null)
                    .build();

            createAndAssert(experiment, apiKey, workspaceName);

            // Creating 5 different experiment ids
            var experimentItems = IntStream
                    .range(0, 2).mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experiment.id())
                            .traceId(traces.get(i).id())
                            .input(traces.get(i).input())
                            .output(traces.get(i).output())
                            .usage(getUsage(spansMap, traces.get(i)))
                            .totalEstimatedCost(getTotalEstimatedCost(spansMap, traces.get(i)))
                            .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(
                                    traces.get(i).startTime(), traces.get(i).endTime()))
                            .datasetItemId(datasetItemBatch.items().get(i).id())
                            .comments(null)
                            .feedbackScores(null)
                            .build())
                    .toList();

            createAndAssert(new ExperimentItemsBatch(Set.copyOf(experimentItems)), apiKey, workspaceName);

            var otherExperimentItems = IntStream.range(0, 3)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experiment.id())
                            .usage(null)
                            .totalEstimatedCost(null)
                            .duration(null)
                            .datasetItemId(datasetItemBatch.items().get(i + 2).id())
                            .comments(null)
                            .feedbackScores(null)
                            .output(null)
                            .input(null)
                            .traceVisibilityMode(null)
                            .build())
                    .toList();

            createAndAssert(new ExperimentItemsBatch(Set.copyOf(otherExperimentItems)), apiKey, workspaceName);

            Set<Column> columns = datasetItemBatch.items()
                    .stream()
                    .flatMap(item -> item.data().entrySet().stream())
                    .map(column -> new Column(column.getKey(), Set.of(getType(column)), "data"))
                    .collect(toSet());

            List<DatasetItem> datasetItems = datasetItemBatch.items()
                    .stream()
                    .sorted(Comparator.comparing(DatasetItem::id).reversed())
                    .toList();

            List<ExperimentItem> expectedExperimentItems = new ArrayList<>();

            expectedExperimentItems.addAll(otherExperimentItems.reversed());
            expectedExperimentItems.addAll(experimentItems.reversed());

            assertPageAndContent(datasetId, List.of(experiment.id()), apiKey, workspaceName, expectedExperimentItems,
                    columns, datasetItems);
        }

        private void assertPageAndContent(UUID datasetId, List<UUID> experimentIds, String apiKey, String workspaceName,
                List<ExperimentItem> expectedExperimentItems, Set<Column> columns, List<DatasetItem> datasetItems) {
            var actualPage = datasetResourceClient.getDatasetItemsWithExperimentItems(datasetId, experimentIds, apiKey,
                    workspaceName);

            assertThat(actualPage.page()).isEqualTo(1);
            assertThat(actualPage.size()).isEqualTo(datasetItems.size());
            assertThat(actualPage.total()).isEqualTo(datasetItems.size());
            assertThat(actualPage.columns()).isEqualTo(columns);

            var actualDatasetItems = actualPage.content();

            assertPage(datasetItems, actualPage.content());

            for (var i = 0; i < actualDatasetItems.size(); i++) {
                var actualDatasetItem = actualDatasetItems.get(i);
                var expectedDatasetItem = datasetItems.get(i);
                var expectedExperimentItem = expectedExperimentItems.get(i);

                assertThat(actualDatasetItem.experimentItems())
                        .usingRecursiveComparison()
                        .ignoringFields(IGNORED_FIELDS_LIST)
                        .withComparatorForType(StatsUtils::bigDecimalComparator, BigDecimal.class)
                        .withComparatorForFields(StatsUtils::closeToEpsilonComparator, "duration")
                        .isEqualTo(List.of(expectedExperimentItem));

                for (var j = 0; j < actualDatasetItem.experimentItems().size(); j++) {
                    var actualExperimentItem = assertFeedbackScoresIgnoredFieldsAndSetThemToNull(
                            actualDatasetItem.experimentItems().get(j), USER);

                    assertThat(actualExperimentItem.feedbackScores())
                            .usingRecursiveComparison()
                            .withComparatorForType(BigDecimal::compareTo, BigDecimal.class)
                            .ignoringCollectionOrder()
                            .isEqualTo(expectedExperimentItem.feedbackScores());

                    assertThat(actualExperimentItem.createdAt())
                            .isAfter(expectedExperimentItem.createdAt());
                    assertThat(actualExperimentItem.lastUpdatedAt())
                            .isAfter(expectedExperimentItem.lastUpdatedAt());

                    assertThat(actualExperimentItem.createdBy())
                            .isEqualTo(USER);
                    assertThat(actualExperimentItem.lastUpdatedBy())
                            .isEqualTo(USER);
                }

                assertThat(actualDatasetItem.createdAt()).isAfter(expectedDatasetItem.createdAt());
                assertThat(actualDatasetItem.lastUpdatedAt()).isAfter(expectedDatasetItem.lastUpdatedAt());
            }
        }

        private Span createSpan(Trace trace, String apiKey, String workspaceName) {
            Span span = factory.manufacturePojo(Span.class).toBuilder()
                    .totalEstimatedCost(BigDecimal.valueOf(PodamUtils.getIntegerInRange(0, 10)))
                    .feedbackScores(null)
                    .totalEstimatedCostVersion(null)
                    .type(SpanType.llm)
                    .errorInfo(null)
                    .comments(null)
                    .traceId(trace.id())
                    .projectName(trace.projectName())
                    .build();

            spanResourceClient.createSpan(span, apiKey, workspaceName);

            return span;
        }

        private BigDecimal getTotalEstimatedCost(Map<UUID, List<Span>> spansMap, Trace trace) {
            return Optional.ofNullable(spansMap.get(trace.id()))
                    .stream()
                    .flatMap(List::stream)
                    .map(Span::totalEstimatedCost)
                    .reduce(BigDecimal::add)
                    .filter(v -> v.compareTo(BigDecimal.ZERO) > 0)
                    .orElse(null);
        }

        private static Map<String, Long> getUsage(Map<UUID, List<Span>> spansMap, Trace trace) {
            return Optional.ofNullable(spansMap.get(trace.id()))
                    .map(spans -> StatsUtils.calculateUsage(
                            spans.stream()
                                    .map(it -> it.usage().entrySet()
                                            .stream()
                                            .collect(Collectors.toMap(
                                                    Map.Entry::getKey,
                                                    entry -> entry.getValue().longValue())))
                                    .toList()))
                    .orElse(null);
        }

        @ParameterizedTest
        @MethodSource("com.comet.opik.api.resources.utils.ImageTruncationArgProvider#provideTestArguments")
        void findWithImageTruncation(JsonNode original, JsonNode expected, boolean truncate) {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            // Creating traces with images to be truncated
            var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream()
                    .map(trace -> trace.toBuilder()
                            .input(original)
                            .output(original)
                            .metadata(original)
                            .build())
                    .toList();

            traceResourceClient.batchCreateTraces(traces, apiKey, workspaceName);

            // Creating the dataset
            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            // Creating 5 dataset items for the dataset above
            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .build();
            var datasetItemBatchWithImage = datasetItemBatch.toBuilder()
                    .items(datasetItemBatch.items().stream()
                            .map(item -> item.toBuilder()
                                    .data(ImmutableMap.of("image", original)).build())
                            .toList())
                    .build();

            putAndAssert(datasetItemBatchWithImage, workspaceName, apiKey);

            // Creating 5 different experiment ids
            var experimentIds = IntStream.range(0, 5).mapToObj(__ -> GENERATOR.generate()).toList();
            List<ExperimentItem> experimentItems = IntStream.range(0, 5)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experimentIds.get(i))
                            .traceId(traces.get(i).id())
                            .usage(null)
                            .totalEstimatedCost(null)
                            .duration(DurationUtils.getDurationInMillisWithSubMilliPrecision(
                                    traces.get(i).startTime(), traces.get(i).endTime()))
                            .datasetItemId(datasetItemBatchWithImage.items().get(i).id()).build())
                    .toList();

            var experimentItemsBatch = ExperimentItemsBatch.builder()
                    .experimentItems(Set.copyOf(experimentItems)).build();

            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            List<List<ExperimentItem>> expectedExperimentItems = experimentItems.stream()
                    .map(item -> List.of(item.toBuilder()
                            .input(expected)
                            .output(expected).build()))
                    .toList();
            var expectedDatasetItems = IntStream.range(0, 5).mapToObj(i -> datasetItemBatchWithImage.items().get(i)
                    .toBuilder()
                    .data(ImmutableMap.of("image", expected))
                    .experimentItems(expectedExperimentItems.get(i))
                    .build()).toList().reversed();

            var page = 1;
            var pageSize = 5;
            var experimentIdsQueryParm = JsonUtils.writeValueAsString(experimentIds);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("page", page)
                    .queryParam("size", pageSize)
                    .queryParam("experiment_ids", experimentIdsQueryParm)
                    .queryParam("truncate", truncate)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_OK);
                var actualPage = actualResponse.readEntity(DatasetItemPage.class);

                assertThat(actualPage.page()).isEqualTo(page);
                assertThat(actualPage.size()).isEqualTo(expectedDatasetItems.size());
                assertThat(actualPage.total()).isEqualTo(expectedDatasetItems.size());

                assertPage(expectedDatasetItems, actualPage.content());

                assertThat(actualPage.content()).hasSize(expectedDatasetItems.size());
                for (int i = 0; i < expectedExperimentItems.size(); i++) {
                    assertThat(actualPage.content().get(i).experimentItems())
                            .usingRecursiveComparison()
                            .withComparatorForType(StatsUtils::bigDecimalComparator, BigDecimal.class)
                            .withComparatorForFields(StatsUtils::closeToEpsilonComparator, "duration")
                            .ignoringFields(IGNORED_FIELDS_LIST)
                            .isEqualTo(expectedExperimentItems.reversed().get(i));
                }
            }
        }

        @ParameterizedTest
        @ValueSource(strings = {"[wrong_payload]", "[0191377d-06ee-7026-8f63-cc5309d1f54b]"})
        void findInvalidExperimentIds(String experimentIds) {
            var expectedErrorMessage = new io.dropwizard.jersey.errors.ErrorMessage(
                    400, "Invalid query param ids '%s'".formatted(experimentIds));

            var datasetId = GENERATOR.generate();
            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("experiment_ids", experimentIds)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(400);

                var actualErrorMessage = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class);

                assertThat(actualErrorMessage).isEqualTo(expectedErrorMessage);
            }
        }

        @Test
        void find__whenNoMatchFound__thenReturnEmptyPageWithAlColumnsFromDataset() {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            List<DatasetItem> items = PodamFactoryUtils.manufacturePojoList(factory, DatasetItem.class);

            var batch = DatasetItemBatch.builder()
                    .items(items)
                    .datasetId(datasetId)
                    .build();
            putAndAssert(batch, workspaceName, apiKey);

            String projectName = RandomStringUtils.randomAlphanumeric(20);
            List<Trace> traces = new ArrayList<>();
            createTraces(items, projectName, workspaceName, apiKey, traces);

            UUID experimentId = GENERATOR.generate();

            List<FeedbackScoreBatchItem> scores = new ArrayList<>();
            createScores(traces, projectName, scores);
            createScoreAndAssert(FeedbackScoreBatch.builder().scores(scores).build(), apiKey, workspaceName);

            List<ExperimentItem> experimentItems = new ArrayList<>();
            createExperimentItems(items, traces, scores, experimentId, experimentItems);

            createAndAssert(
                    ExperimentItemsBatch.builder()
                            .experimentItems(Set.copyOf(experimentItems))
                            .build(),
                    apiKey,
                    workspaceName);

            Set<Column> columns = getColumns(items.stream().map(DatasetItem::data).toList());

            List<Filter> filters = List.of(ExperimentsComparisonFilter.builder()
                    .type(FieldType.STRING)
                    .value(RandomStringUtils.randomAlphanumeric(16))
                    .field(RandomStringUtils.randomAlphanumeric(22))
                    .operator(Operator.EQUAL)
                    .build());

            assertDatasetExperimentPage(datasetId, experimentId, filters, apiKey, workspaceName, columns, List.of());
        }

        @ParameterizedTest
        @ValueSource(strings = {"$..test", "$meta_field", "[", "]", "[..]",})
        void findInvalidJsonPaths(String path) {

            Dataset dataset = Dataset.builder().name(GENERATOR.generate().toString()).build();

            var datasetId = createAndAssert(dataset);

            Experiment experiment = Experiment.builder()
                    .id(GENERATOR.generate())
                    .name(GENERATOR.generate().toString())
                    .datasetName(dataset.name())
                    .build();

            createAndAssert(experiment, API_KEY, TEST_WORKSPACE);

            var field = RandomStringUtils.randomAlphanumeric(5);

            var filters = List
                    .of(new ExperimentsComparisonFilter(field, FieldType.DICTIONARY, Operator.EQUAL, path, "10"));
            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("experiment_ids", JsonUtils.writeValueAsString(List.of(experiment.id())))
                    .queryParam("filters", toURLEncodedQueryParam(filters))
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, API_KEY)
                    .header(WORKSPACE_HEADER, TEST_WORKSPACE)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);

                var datasetItemPage = actualResponse.readEntity(DatasetItemPage.class);

                assertThat(datasetItemPage.content()).isEmpty();
                assertThat(datasetItemPage.total()).isZero();
                assertThat(datasetItemPage.size()).isZero();
                assertThat(datasetItemPage.page()).isOne();
            }
        }

        @ParameterizedTest
        @MethodSource
        void find__whenFilteringBySupportedFields__thenReturnMatchingRows(Filter filter) {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            List<DatasetItem> datasetItems = new ArrayList<>();
            createDatasetItems(datasetItems);
            var batch = DatasetItemBatch.builder()
                    .items(datasetItems)
                    .datasetId(datasetId)
                    .build();
            putAndAssert(batch, workspaceName, apiKey);

            String projectName = RandomStringUtils.secure().nextAlphabetic(20);
            List<Trace> traces = new ArrayList<>();
            createTraces(datasetItems, projectName, workspaceName, apiKey, traces);

            UUID experimentId = GENERATOR.generate();

            List<FeedbackScoreBatchItem> scores = new ArrayList<>();
            createScores(traces, projectName, scores);
            createScoreAndAssert(FeedbackScoreBatch.builder().scores(scores).build(), apiKey, workspaceName);

            List<ExperimentItem> experimentItems = new ArrayList<>();
            createExperimentItems(datasetItems, traces, scores, experimentId, experimentItems);

            createAndAssert(
                    ExperimentItemsBatch.builder()
                            .experimentItems(Set.copyOf(experimentItems))
                            .build(),
                    apiKey,
                    workspaceName);

            Set<Column> columns = getColumns(datasetItems.stream().map(DatasetItem::data).toList());

            List<Filter> filters = List.of(filter);

            var expectedDatasetItems = filter.operator() != Operator.NOT_EQUAL
                    ? List.of(datasetItems.getFirst())
                    : datasetItems.subList(1, datasetItems.size()).reversed();
            var expectedExperimentItems = filter.operator() != Operator.NOT_EQUAL
                    ? List.of(experimentItems.getFirst())
                    : experimentItems.subList(1, experimentItems.size()).reversed();

            var actualPage = assertDatasetExperimentPage(datasetId, experimentId, filters, apiKey, workspaceName,
                    columns, expectedDatasetItems);

            assertDatasetItemExperiments(actualPage, expectedDatasetItems, expectedExperimentItems);
        }

        Stream<Arguments> find__whenFilteringBySupportedFields__thenReturnMatchingRows() {
            return Stream.of(
                    arguments(new ExperimentsComparisonFilter("sql_tag",
                            FieldType.STRING, Operator.EQUAL, null, "sql_test")),
                    arguments(new ExperimentsComparisonFilter("sql_tag",
                            FieldType.STRING, Operator.NOT_EQUAL, null, "sql_test")),
                    arguments(new ExperimentsComparisonFilter("sql_tag",
                            FieldType.STRING, Operator.CONTAINS, null, "sql_")),
                    arguments(new ExperimentsComparisonFilter("json_node",
                            FieldType.DICTIONARY, Operator.EQUAL, "test2", "12338")),
                    arguments(new ExperimentsComparisonFilter("json_node",
                            FieldType.DICTIONARY, Operator.NOT_EQUAL, "test2", "12338")),
                    arguments(new ExperimentsComparisonFilter("sql_rate",
                            FieldType.NUMBER, Operator.LESS_THAN, null, "101")),
                    arguments(new ExperimentsComparisonFilter("sql_rate",
                            FieldType.NUMBER, Operator.GREATER_THAN, null, "99")),
                    arguments(new ExperimentsComparisonFilter("feedback_scores",
                            FieldType.FEEDBACK_SCORES_NUMBER, Operator.EQUAL, "sql_cost", "10")),
                    arguments(new ExperimentsComparisonFilter("feedback_scores",
                            FieldType.FEEDBACK_SCORES_NUMBER, Operator.NOT_EQUAL, "sql_cost", "10")),
                    arguments(new ExperimentsComparisonFilter("output",
                            FieldType.STRING, Operator.CONTAINS, null, "sql_cost")),
                    arguments(new ExperimentsComparisonFilter("expected_output",
                            FieldType.DICTIONARY, Operator.CONTAINS, "output", "sql_cost")),
                    arguments(new ExperimentsComparisonFilter("metadata",
                            FieldType.DICTIONARY, Operator.EQUAL, "sql_cost", "10")),
                    arguments(new ExperimentsComparisonFilter("meta_field",
                            FieldType.DICTIONARY, Operator.CONTAINS, "version[*]", "10")),
                    arguments(new ExperimentsComparisonFilter("releases",
                            FieldType.DICTIONARY, Operator.CONTAINS, "$[1].version", "1.1")),
                    arguments(new ExperimentsComparisonFilter("meta_field",
                            FieldType.DICTIONARY, Operator.CONTAINS, ".version[1]", "11")));

        }

        @Test
        void findWithDeletedDatasetItems() {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            // Creating a trace with input and output
            var trace = factory.manufacturePojo(Trace.class);
            createAndAssert(trace, workspaceName, apiKey);

            // Creating the dataset
            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            // Creating a dataset item
            var datasetItem = factory.manufacturePojo(DatasetItem.class);
            var datasetItemBatch = DatasetItemBatch.builder()
                    .datasetId(datasetId)
                    .items(List.of(datasetItem))
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            // Creating experiment and experiment item
            var experimentId = GENERATOR.generate();
            var experimentItem = factory.manufacturePojo(ExperimentItem.class).toBuilder()
                    .experimentId(experimentId)
                    .datasetItemId(datasetItem.id())
                    .traceId(trace.id())
                    .input(trace.input())
                    .output(trace.output())
                    .build();

            var experimentItemsBatch = ExperimentItemsBatch.builder()
                    .experimentItems(Set.of(experimentItem))
                    .build();
            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Delete the dataset item using the correct POST endpoint
            var deleteRequest = new DatasetItemsDelete(List.of(datasetItem.id()));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path("items")
                    .path("delete")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .post(Entity.json(deleteRequest))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
            }

            // Create expected dataset item with null source and empty data since it was deleted
            var expectedDatasetItem = datasetItem.toBuilder()
                    .source(null)
                    .data(Map.of()) // Data should be empty when dataset item is deleted
                    .traceId(null) // NULL because dataset item was hard-deleted
                    .spanId(null) // NULL because dataset item was hard-deleted
                    .build();

            // Create expected experiment item with adjusted fields to match API response
            var expectedExperimentItem = experimentItem.toBuilder()
                    .feedbackScores(null) // API returns null for feedback scores in this context
                    .comments(null) // API returns null for comments in this context
                    .totalEstimatedCost(null) // API returns null for totalEstimatedCost in this context
                    .usage(null) // API returns null for usage in this context
                    // Don't set duration - let it be compared as-is from the API response
                    .build();

            // Query for dataset items with experiment items using assertion helper
            var actualPage = assertDatasetExperimentPage(datasetId, experimentId, List.of(),
                    apiKey, workspaceName, Set.of(), List.of(expectedDatasetItem));

            // Get the actual duration from the API response to use in expected experiment item
            var actualExperimentItem = actualPage.content().get(0).experimentItems().get(0);
            var expectedExperimentItemWithActualDuration = expectedExperimentItem.toBuilder()
                    .duration(actualExperimentItem.duration()) // Use actual duration from API
                    .build();

            // Verify the experiment items are properly associated using assertion helper
            assertDatasetItemExperiments(actualPage, List.of(expectedDatasetItem),
                    List.of(expectedExperimentItemWithActualDuration));
        }

        private void createExperimentItems(List<DatasetItem> items, List<Trace> traces,
                List<FeedbackScoreBatchItem> scores, UUID experimentId, List<ExperimentItem> experimentItems) {
            for (int i = 0; i < items.size(); i++) {
                var item = items.get(i);
                var trace = traces.get(i);
                var score = scores.get(i);

                var experimentItem = ExperimentItem.builder()
                        .id(GENERATOR.generate())
                        .datasetItemId(item.id())
                        .traceId(trace.id())
                        .experimentId(experimentId)
                        .input(trace.input())
                        .output(trace.output())
                        .traceVisibilityMode(VisibilityMode.DEFAULT)
                        .feedbackScores(score == null
                                ? null
                                : Stream.of(score)
                                        .map(FeedbackScoreMapper.INSTANCE::toFeedbackScore)
                                        .toList())
                        .build();

                experimentItems.add(experimentItem);
            }
        }

        private void createScores(List<Trace> traces, String projectName, List<FeedbackScoreBatchItem> scores) {
            for (int i = 0; i < traces.size(); i++) {
                var trace = traces.get(i);

                var score = FeedbackScoreBatchItem.builder()
                        .name("sql_cost")
                        .value(BigDecimal.valueOf(i == 0 ? 10 : i))
                        .source(ScoreSource.SDK)
                        .id(trace.id())
                        .projectName(projectName)
                        .build();

                scores.add(score);
            }
        }

        private void createTraces(List<DatasetItem> items, String projectName, String workspaceName, String apiKey,
                List<Trace> traces) {
            for (int i = 0; i < items.size(); i++) {
                var item = items.get(i);
                var trace = Trace.builder()
                        .id(GENERATOR.generate())
                        .input(item.data().get("input"))
                        .output(item.data().get("expected_output"))
                        .projectName(projectName)
                        .startTime(Instant.now())
                        .name("trace-" + i)
                        .build();

                createAndAssert(trace, workspaceName, apiKey);
                traces.add(trace);
            }
        }

        private void createDatasetItems(List<DatasetItem> items) {
            for (int i = 0; i < 5; i++) {
                if (i == 0) {
                    DatasetItem item = factory.manufacturePojo(DatasetItem.class)
                            .toBuilder()
                            .source(DatasetItemSource.SDK)
                            .data(new HashMap<>() {
                                {
                                    put("sql_tag", JsonUtils.readTree("sql_test"));
                                    put("sql_rate", JsonUtils.readTree(100));
                                    put("input", JsonUtils
                                            .getJsonNodeFromString(
                                                    JsonUtils.writeValueAsString(Map.of("input", "sql_cost"))));
                                    put("expected_output", JsonUtils
                                            .getJsonNodeFromString(
                                                    JsonUtils.writeValueAsString(Map.of("output", "sql_cost"))));
                                    put("metadata", JsonUtils
                                            .getJsonNodeFromString(
                                                    JsonUtils.writeValueAsString(Map.of("sql_cost", 10))));
                                    put("meta_field",
                                            JsonUtils.readTree(Map.of("version", new String[]{"10", "11", "12"})));
                                    put("releases", JsonUtils.readTree(
                                            List.of(
                                                    Map.of("fixes", new String[]{"10", "11", "12"}, "version", "1.0"),
                                                    Map.of("fixes", new String[]{"10", "11", "12"}, "version", "1.1"),
                                                    Map.of("fixes", new String[]{"10", "45", "30"}, "version",
                                                            "1.2"))));
                                    put("json_node", JsonUtils.readTree(Map.of("test", "1233", "test2", "12338")));
                                    put(RandomStringUtils.randomAlphanumeric(5),
                                            BigIntegerNode.valueOf(new BigInteger("18446744073709551615")));
                                    put(RandomStringUtils.randomAlphanumeric(5), DoubleNode.valueOf(132432432.79995));
                                    put(RandomStringUtils.randomAlphanumeric(5),
                                            DoubleNode.valueOf(1.1844674407370955444555));
                                    put(RandomStringUtils.randomAlphanumeric(5), IntNode.valueOf(100000000));
                                    put(RandomStringUtils.randomAlphanumeric(5), BooleanNode.valueOf(true));
                                }
                            })
                            .traceId(null)
                            .spanId(null)
                            .build();

                    items.add(item);
                } else {
                    var item = factory.manufacturePojo(DatasetItem.class);

                    items.add(item);
                }
            }
        }

        private void createScoreAndAssert(FeedbackScoreBatchContainer feedbackScoreBatch, String apiKey,
                String workspaceName) {
            try (var actualResponse = client.target(getTracesPath())
                    .path("feedback-scores")
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .put(Entity.json(feedbackScoreBatch))) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
                assertThat(actualResponse.hasEntity()).isFalse();
            }
        }

        @Test
        void find__whenFilteringFeedbackScoresEmpty__thenReturnMatchingRows() {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            List<DatasetItem> datasetItems = new ArrayList<>();
            createDatasetItems(datasetItems);
            var batch = DatasetItemBatch.builder()
                    .items(datasetItems)
                    .datasetId(datasetId)
                    .build();
            putAndAssert(batch, workspaceName, apiKey);

            String projectName = RandomStringUtils.randomAlphanumeric(20);
            List<Trace> traces = new ArrayList<>();
            createTraces(datasetItems, projectName, workspaceName, apiKey, traces);

            UUID experimentId = GENERATOR.generate();

            List<FeedbackScoreBatchItem> scores = new ArrayList<>();
            createScores(traces.subList(0, traces.size() - 1), projectName, scores);
            scores = scores.stream()
                    .map(score -> score.toBuilder()
                            .name(factory.manufacturePojo(String.class))
                            .build())
                    .collect(toList());
            createScoreAndAssert(FeedbackScoreBatch.builder().scores(scores).build(), apiKey, workspaceName);

            List<ExperimentItem> experimentItems = new ArrayList<>();
            createExperimentItems(datasetItems, traces, Stream.concat(scores.stream(),
                    Stream.of((FeedbackScoreBatchItem) null)).toList(),
                    experimentId, experimentItems);

            createAndAssert(
                    ExperimentItemsBatch.builder()
                            .experimentItems(Set.copyOf(experimentItems))
                            .build(),
                    apiKey,
                    workspaceName);

            Set<Column> columns = getColumns(datasetItems.stream().map(DatasetItem::data).toList());

            var isNotEmptyFilter = List.of(
                    ExperimentsComparisonFilter.builder()
                            .field(ExperimentsComparisonValidKnownField.FEEDBACK_SCORES.getQueryParamField())
                            .operator(Operator.IS_NOT_EMPTY)
                            .key(scores.getFirst().name())
                            .value("")
                            .build());

            var actualPageIsNotEmpty = assertDatasetExperimentPage(datasetId, experimentId, isNotEmptyFilter, apiKey,
                    workspaceName, columns, List.of(datasetItems.getFirst()));

            assertDatasetItemExperiments(actualPageIsNotEmpty, List.of(datasetItems.getFirst()),
                    List.of(experimentItems.getFirst()));

            var isEmptyFilter = List.of(
                    ExperimentsComparisonFilter.builder()
                            .field(ExperimentsComparisonValidKnownField.FEEDBACK_SCORES.getQueryParamField())
                            .operator(Operator.IS_EMPTY)
                            .key(scores.getFirst().name())
                            .value("")
                            .build());

            var actualPageIsEmpty = assertDatasetExperimentPage(datasetId, experimentId, isEmptyFilter, apiKey,
                    workspaceName, columns, datasetItems.subList(1, datasetItems.size()).reversed());

            assertDatasetItemExperiments(actualPageIsEmpty, datasetItems.subList(1, datasetItems.size()).reversed(),
                    experimentItems.subList(1, datasetItems.size()).reversed());
        }

        @ParameterizedTest
        @MethodSource
        void find__whenFilterInvalidOperatorForFieldType__thenReturn400(ExperimentsComparisonFilter filter) {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var expectedError = new io.dropwizard.jersey.errors.ErrorMessage(
                    400,
                    "Invalid operator '%s' for field '%s' of type '%s'".formatted(
                            filter.operator().getQueryParamOperator(),
                            filter.field().getQueryParamField(),
                            filter.field().getType()));

            var datasetId = GENERATOR.generate();
            var experimentId = GENERATOR.generate();
            var filters = List.of(filter);

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("experiment_ids", JsonUtils.writeValueAsString(List.of(experimentId)))
                    .queryParam("filters", toURLEncodedQueryParam(filters))
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(400);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class);
                assertThat(actualError).isEqualTo(expectedError);
            }

        }

        static Stream<Arguments> find__whenFilterInvalidOperatorForFieldType__thenReturn400() {
            return Stream.of(
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("feedback_scores")
                            .type(FieldType.FEEDBACK_SCORES_NUMBER)
                            .operator(Operator.CONTAINS)
                            .value(RandomStringUtils.randomAlphanumeric(10))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("feedback_scores")
                            .type(FieldType.FEEDBACK_SCORES_NUMBER)
                            .operator(Operator.NOT_CONTAINS)
                            .value(RandomStringUtils.randomAlphanumeric(10))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("feedback_scores")
                            .type(FieldType.FEEDBACK_SCORES_NUMBER)
                            .operator(Operator.STARTS_WITH)
                            .value(RandomStringUtils.randomAlphanumeric(10))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("feedback_scores")
                            .type(FieldType.FEEDBACK_SCORES_NUMBER)
                            .operator(Operator.ENDS_WITH)
                            .value(RandomStringUtils.randomAlphanumeric(10))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("input")
                            .type(FieldType.STRING)
                            .operator(Operator.GREATER_THAN_EQUAL)
                            .value(RandomStringUtils.randomNumeric(3))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("input")
                            .type(FieldType.STRING)
                            .operator(Operator.LESS_THAN_EQUAL)
                            .value(RandomStringUtils.randomNumeric(3))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("output")
                            .type(FieldType.STRING)
                            .operator(Operator.GREATER_THAN_EQUAL)
                            .value(RandomStringUtils.randomNumeric(3))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("output")
                            .type(FieldType.STRING)
                            .operator(Operator.LESS_THAN_EQUAL)
                            .value(RandomStringUtils.randomNumeric(3))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("expected_output")
                            .type(FieldType.STRING)
                            .operator(Operator.GREATER_THAN_EQUAL)
                            .value(RandomStringUtils.randomNumeric(3))
                            .build()),
                    Arguments.of(ExperimentsComparisonFilter.builder()
                            .field("expected_output")
                            .type(FieldType.STRING)
                            .operator(Operator.LESS_THAN_EQUAL)
                            .value(RandomStringUtils.randomNumeric(3))
                            .build()));
        }

        @ParameterizedTest
        @MethodSource
        void find__whenFilterInvalidFieldTypeForDynamicFields__thenReturn400(String filters) {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var expectedError = new io.dropwizard.jersey.errors.ErrorMessage(
                    400,
                    "Invalid filters query parameter '%s'".formatted(filters));

            var datasetId = GENERATOR.generate();
            var experimentId = GENERATOR.generate();

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("experiment_ids", JsonUtils.writeValueAsString(List.of(experimentId)))
                    .queryParam("filters", URLEncoder.encode(filters, StandardCharsets.UTF_8))
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(400);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualError = actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class);
                assertThat(actualError).isEqualTo(expectedError);
            }

        }

        static Stream<Arguments> find__whenFilterInvalidFieldTypeForDynamicFields__thenReturn400() {
            String template = "[{\"field\":\"%s\",\"type\":\"%s\",\"operator\":\"%s\",\"value\":\"%s\"}]";
            String field = RandomStringUtils.randomAlphanumeric(5);
            return Stream.of(
                    Arguments.of(template.formatted(
                            field,
                            FieldType.FEEDBACK_SCORES_NUMBER.getQueryParamType(),
                            Operator.EQUAL.getQueryParamOperator(),
                            RandomStringUtils.randomAlphanumeric(10))),
                    Arguments.of(template.formatted(
                            field,
                            FieldType.DATE_TIME.getQueryParamType(),
                            Operator.EQUAL.getQueryParamOperator(),
                            Instant.now().toString())),
                    Arguments.of(template.formatted(
                            field,
                            FieldType.LIST.getQueryParamType(),
                            Operator.CONTAINS.getQueryParamOperator(),
                            RandomStringUtils.randomAlphanumeric(10))));
        }
    }

    private DatasetItemPage assertDatasetExperimentPage(UUID datasetId, UUID experimentId,
            List<? extends Filter> filters,
            String apiKey, String workspaceName, Set<Column> columns, List<DatasetItem> datasetItems) {
        try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                .path(datasetId.toString())
                .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                .queryParam("experiment_ids", JsonUtils.writeValueAsString(List.of(experimentId)))
                .queryParam("filters", toURLEncodedQueryParam(filters))
                .request()
                .header(HttpHeaders.AUTHORIZATION, apiKey)
                .header(WORKSPACE_HEADER, workspaceName)
                .get()) {

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
            assertThat(actualResponse.hasEntity()).isTrue();

            var actualPage = actualResponse.readEntity(DatasetItemPage.class);

            assertDatasetItemPage(actualPage, datasetItems, columns, 1);

            return actualPage;
        }
    }

    @DisplayName("Find dataset items with experiment items sorting test")
    @Nested
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class FindDatasetItemsWithExperimentItemsSortingTest {

        @ParameterizedTest
        @MethodSource
        @DisplayName("when sorting dataset items with experiment items, then return items sorted by valid fields")
        void findDatasetItemsWithExperimentItems__whenSorting__thenReturnItemsSortedByValidFields(
                SortingField sorting) {

            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            // Create dataset
            var dataset = factory.manufacturePojo(Dataset.class).toBuilder().id(null).build();
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            // Create project and traces
            String projectName = GENERATOR.generate().toString();
            List<Trace> traces = IntStream.range(0, 5)
                    .mapToObj(i -> factory.manufacturePojo(Trace.class).toBuilder()
                            .projectName(projectName)
                            .build())
                    .toList();

            traces.forEach(trace -> createAndAssert(trace, workspaceName, apiKey));

            // Create dataset items
            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .items(traces.stream()
                            .map(trace -> factory.manufacturePojo(DatasetItem.class).toBuilder()
                                    .datasetId(datasetId)
                                    .traceId(trace.id())
                                    .spanId(null)
                                    .source(DatasetItemSource.TRACE)
                                    .build())
                            .toList())
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            // Create experiment
            var experiment = factory.manufacturePojo(Experiment.class).toBuilder()
                    .datasetName(dataset.name())
                    .promptVersion(null)
                    .promptVersions(null)
                    .build();

            createAndAssert(experiment, apiKey, workspaceName);

            // Create experiment items
            var experimentItems = IntStream.range(0, traces.size())
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experiment.id())
                            .datasetItemId(datasetItemBatch.items().get(i).id())
                            .traceId(traces.get(i).id())
                            .build())
                    .collect(Collectors.toSet());

            createAndAssert(new ExperimentItemsBatch(experimentItems), apiKey, workspaceName);

            // Fetch dataset items with sorting (first call)
            var sortingParam = URLEncoder.encode(JsonUtils.writeValueAsString(List.of(sorting)),
                    StandardCharsets.UTF_8);
            var experimentIdsParam = JsonUtils.writeValueAsString(List.of(experiment.id()));

            List<UUID> firstFetchIds;
            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("page", 1)
                    .queryParam("size", 10)
                    .queryParam("experiment_ids", experimentIdsParam)
                    .queryParam("sorting", sortingParam)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualPage = actualResponse.readEntity(DatasetItemPage.class);
                assertThat(actualPage.content()).isNotEmpty();
                assertThat(actualPage.content()).hasSize(5);

                firstFetchIds = actualPage.content().stream().map(DatasetItem::id).toList();
            }

            // Fetch again to verify deterministic ordering (second call)
            List<UUID> secondFetchIds;
            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("page", 1)
                    .queryParam("size", 10)
                    .queryParam("experiment_ids", experimentIdsParam)
                    .queryParam("sorting", sortingParam)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualPage = actualResponse.readEntity(DatasetItemPage.class);
                assertThat(actualPage.content()).isNotEmpty();

                secondFetchIds = actualPage.content().stream().map(DatasetItem::id).toList();
            }

            // Verify sorting is deterministic - both fetches should return the same order
            assertThat(firstFetchIds).containsExactlyElementsOf(secondFetchIds);
        }

        private Stream<Arguments> findDatasetItemsWithExperimentItems__whenSorting__thenReturnItemsSortedByValidFields() {
            return Stream.of(
                    // ID field sorting
                    Arguments.of(SortingField.builder().field(SortableFields.ID).direction(Direction.ASC).build()),
                    Arguments.of(SortingField.builder().field(SortableFields.ID).direction(Direction.DESC).build()),

                    // CREATED_AT field sorting
                    Arguments.of(
                            SortingField.builder().field(SortableFields.CREATED_AT).direction(Direction.ASC).build()),
                    Arguments.of(
                            SortingField.builder().field(SortableFields.CREATED_AT).direction(Direction.DESC).build()),

                    // LAST_UPDATED_AT field sorting
                    Arguments.of(SortingField.builder().field(SortableFields.LAST_UPDATED_AT).direction(Direction.ASC)
                            .build()),
                    Arguments.of(SortingField.builder().field(SortableFields.LAST_UPDATED_AT).direction(Direction.DESC)
                            .build()),

                    // CREATED_BY field sorting
                    Arguments.of(
                            SortingField.builder().field(SortableFields.CREATED_BY).direction(Direction.ASC).build()),
                    Arguments.of(
                            SortingField.builder().field(SortableFields.CREATED_BY).direction(Direction.DESC).build()),

                    // LAST_UPDATED_BY field sorting
                    Arguments.of(SortingField.builder().field(SortableFields.LAST_UPDATED_BY).direction(Direction.ASC)
                            .build()),
                    Arguments.of(SortingField.builder().field(SortableFields.LAST_UPDATED_BY).direction(Direction.DESC)
                            .build()),

                    // Dynamic dataset fields (data.*) - Using explicit prefix notation
                    // Users must explicitly use "data." prefix for dataset columns
                    Arguments.of(SortingField.builder().field("data.expected_answer").direction(Direction.ASC).build()),
                    Arguments
                            .of(SortingField.builder().field("data.expected_answer").direction(Direction.DESC).build()),
                    Arguments.of(SortingField.builder().field("data.input").direction(Direction.ASC).build()),
                    Arguments.of(SortingField.builder().field("data.input").direction(Direction.DESC).build()),
                    // Output fields
                    Arguments.of(SortingField.builder().field("output.translation").direction(Direction.ASC).build()),
                    Arguments
                            .of(SortingField.builder().field("output.translation").direction(Direction.DESC).build()),
                    Arguments.of(SortingField.builder().field("output.quality_score").direction(Direction.ASC).build()),
                    Arguments
                            .of(SortingField.builder().field("output.quality_score").direction(Direction.DESC).build()),
                    // Input fields
                    Arguments.of(SortingField.builder().field("input.text").direction(Direction.ASC).build()),
                    Arguments.of(SortingField.builder().field("input.text").direction(Direction.DESC).build()),
                    Arguments.of(SortingField.builder().field("input.language").direction(Direction.ASC).build()),
                    Arguments.of(SortingField.builder().field("input.language").direction(Direction.DESC).build()),
                    // Metadata fields
                    Arguments.of(SortingField.builder().field("metadata.task_type").direction(Direction.ASC).build()),
                    Arguments
                            .of(SortingField.builder().field("metadata.task_type").direction(Direction.DESC).build()),
                    Arguments.of(SortingField.builder().field("metadata.model").direction(Direction.ASC).build()),
                    Arguments.of(SortingField.builder().field("metadata.model").direction(Direction.DESC).build()));
        }

        @Test
        @DisplayName("when sorting with invalid field, then ignore and return success")
        void findDatasetItemsWithExperimentItems__whenSortingWithInvalidField__thenIgnoreAndReturnSuccess() {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder().id(null).build();
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            String projectName = GENERATOR.generate().toString();
            var trace = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(projectName)
                    .build();
            createAndAssert(trace, workspaceName, apiKey);

            var datasetItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .datasetId(datasetId)
                    .traceId(trace.id())
                    .spanId(null)
                    .source(DatasetItemSource.TRACE)
                    .build();

            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .items(List.of(datasetItem))
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            var experiment = factory.manufacturePojo(Experiment.class).toBuilder()
                    .datasetName(dataset.name())
                    .promptVersion(null)
                    .promptVersions(null)
                    .build();

            createAndAssert(experiment, apiKey, workspaceName);

            var experimentItem = factory.manufacturePojo(ExperimentItem.class).toBuilder()
                    .experimentId(experiment.id())
                    .datasetItemId(datasetItem.id())
                    .traceId(trace.id())
                    .build();

            createAndAssert(new ExperimentItemsBatch(Set.of(experimentItem)), apiKey, workspaceName);

            // Use invalid sorting field - should be gracefully ignored
            var invalidSorting = SortingField.builder()
                    .field("invalid_field")
                    .direction(Direction.ASC)
                    .build();
            var sortingParam = URLEncoder.encode(JsonUtils.writeValueAsString(List.of(invalidSorting)),
                    StandardCharsets.UTF_8);
            var experimentIdsParam = JsonUtils.writeValueAsString(List.of(experiment.id()));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("page", 1)
                    .queryParam("size", 10)
                    .queryParam("experiment_ids", experimentIdsParam)
                    .queryParam("sorting", sortingParam)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                // Verify graceful degradation: request succeeds with invalid sorting
                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualPage = actualResponse.readEntity(DatasetItemPage.class);
                assertThat(actualPage).isNotNull();
                // Verify data is returned despite invalid sorting field
                assertThat(actualPage.content()).isNotEmpty();
                assertThat(actualPage.content()).hasSize(1);
            }
        }

        @Test
        @DisplayName("when sorting without sorting parameter, then return unsorted items")
        void findDatasetItemsWithExperimentItems__whenNoSortingParameter__thenReturnUnsortedItems() {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder().id(null).build();
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            String projectName = GENERATOR.generate().toString();
            var trace = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(projectName)
                    .build();
            createAndAssert(trace, workspaceName, apiKey);

            var datasetItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .datasetId(datasetId)
                    .traceId(trace.id())
                    .spanId(null)
                    .source(DatasetItemSource.TRACE)
                    .build();

            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .items(List.of(datasetItem))
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            var experiment = factory.manufacturePojo(Experiment.class).toBuilder()
                    .datasetName(dataset.name())
                    .promptVersion(null)
                    .promptVersions(null)
                    .build();

            createAndAssert(experiment, apiKey, workspaceName);

            var experimentItem = factory.manufacturePojo(ExperimentItem.class).toBuilder()
                    .experimentId(experiment.id())
                    .datasetItemId(datasetItem.id())
                    .traceId(trace.id())
                    .build();

            createAndAssert(new ExperimentItemsBatch(Set.of(experimentItem)), apiKey, workspaceName);

            // Fetch without sorting parameter
            var experimentIdsParam = JsonUtils.writeValueAsString(List.of(experiment.id()));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("page", 1)
                    .queryParam("size", 10)
                    .queryParam("experiment_ids", experimentIdsParam)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualPage = actualResponse.readEntity(DatasetItemPage.class);
                assertThat(actualPage.content()).isNotEmpty();
            }
        }

        @Test
        @DisplayName("when workspace allows dynamic sorting, then sortableBy fields are present in response")
        void findDatasetItemsWithExperimentItems__whenWorkspaceAllowsDynamicSorting__thenSortableByFieldsPresent() {
            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder().id(null).build();
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            String projectName = GENERATOR.generate().toString();
            var trace = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(projectName)
                    .build();
            createAndAssert(trace, workspaceName, apiKey);

            var datasetItem = factory.manufacturePojo(DatasetItem.class).toBuilder()
                    .datasetId(datasetId)
                    .traceId(trace.id())
                    .spanId(null)
                    .source(DatasetItemSource.TRACE)
                    .build();

            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .items(List.of(datasetItem))
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            var experiment = factory.manufacturePojo(Experiment.class).toBuilder()
                    .datasetName(dataset.name())
                    .promptVersion(null)
                    .promptVersions(null)
                    .build();

            createAndAssert(experiment, apiKey, workspaceName);

            var experimentItem = factory.manufacturePojo(ExperimentItem.class).toBuilder()
                    .experimentId(experiment.id())
                    .datasetItemId(datasetItem.id())
                    .traceId(trace.id())
                    .build();

            createAndAssert(new ExperimentItemsBatch(Set.of(experimentItem)), apiKey, workspaceName);

            // Test configuration has maxSizeToAllowSorting: -1 (no limit), so dynamic sorting is enabled
            var experimentIdsParam = JsonUtils.writeValueAsString(List.of(experiment.id()));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("page", 1)
                    .queryParam("size", 10)
                    .queryParam("experiment_ids", experimentIdsParam)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualPage = actualResponse.readEntity(DatasetItemPage.class);
                assertThat(actualPage.content()).isNotEmpty();

                // Verify sortableBy fields are present (workspace allows dynamic sorting)
                assertThat(actualPage.sortableBy()).isNotNull();
                assertThat(actualPage.sortableBy()).isNotEmpty();

                // Verify it contains expected sortable fields
                assertThat(actualPage.sortableBy()).contains(
                        "id",
                        "created_at",
                        "last_updated_at",
                        "created_by",
                        "last_updated_by",
                        "data.*",
                        "output.*",
                        "input.*",
                        "metadata.*",
                        "duration",
                        "feedback_scores.*",
                        "comments");
            }
        }

        @Test
        @DisplayName("when workspace size exceeds limit, then dynamic sorting is disabled and sortableBy is empty")
        void findDatasetItemsWithExperimentItems__whenWorkspaceExceedsSize__thenDynamicSortingDisabled() {
            /*
             * Note: This test verifies the workspace size protection mechanism.
             *
             * The protection works as follows:
             * 1. WorkspaceMetadataService calculates workspace size based on spans data
             * 2. If workspace size exceeds maxSizeToAllowSorting, dynamic sorting is disabled
             * 3. The API strips sortableBy fields and ignores sorting parameters
             *
             * Test configuration (config-test.yml):
             *   workspaceSettings.maxSizeToAllowSorting: -1 (unlimited for tests)
             *
             * In production:
             *   - maxSizeToAllowSorting is set to a reasonable limit (e.g., 100 GB)
             *   - Large workspaces automatically disable dynamic sorting
             *   - This prevents expensive ClickHouse queries on large datasets
             *
             * The workspace metadata check is implemented in:
             *   - ScopeMetadata.canUseDynamicSorting()
             *   - DatasetsResource.findDatasetItemsWithExperimentItems() (lines 440-446, 464-470)
             *
             * To test this scenario in a real environment:
             *   1. Set maxSizeToAllowSorting to a low value (e.g., 0.1 GB)
             *   2. Insert enough spans data to exceed the threshold
             *   3. Verify sortableBy is empty in API response
             *   4. Verify sorting parameters are ignored
             *
             * This test documents the protection mechanism. The actual workspace size
             * calculation and threshold logic is tested in WorkspaceMetadataService tests.
             */

            String workspaceName = UUID.randomUUID().toString();
            String apiKey = UUID.randomUUID().toString();
            String workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class).toBuilder().id(null).build();
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            String projectName = GENERATOR.generate().toString();

            // Create multiple traces and spans to generate workspace data
            // This simulates a workspace with meaningful data for size calculation
            List<Trace> traces = IntStream.range(0, 10)
                    .mapToObj(i -> factory.manufacturePojo(Trace.class).toBuilder()
                            .projectName(projectName)
                            .input(JsonUtils.getJsonNodeFromString(
                                    "{\"text\": \"Test input " + i + " with some content\"}"))
                            .output(JsonUtils.getJsonNodeFromString(
                                    "{\"result\": \"Test output " + i + " with some content\"}"))
                            .build())
                    .toList();

            traces.forEach(trace -> createAndAssert(trace, workspaceName, apiKey));

            var datasetItems = traces.stream()
                    .map(trace -> factory.manufacturePojo(DatasetItem.class).toBuilder()
                            .datasetId(datasetId)
                            .traceId(trace.id())
                            .spanId(null)
                            .source(DatasetItemSource.TRACE)
                            .build())
                    .toList();

            var datasetItemBatch = DatasetItemBatch.builder()
                    .datasetId(datasetId)
                    .items(datasetItems)
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            var experiment = factory.manufacturePojo(Experiment.class).toBuilder()
                    .datasetName(dataset.name())
                    .promptVersion(null)
                    .promptVersions(null)
                    .build();

            createAndAssert(experiment, apiKey, workspaceName);

            var experimentItems = IntStream.range(0, traces.size())
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experiment.id())
                            .datasetItemId(datasetItems.get(i).id())
                            .traceId(traces.get(i).id())
                            .input(traces.get(i).input())
                            .output(traces.get(i).output())
                            .build())
                    .collect(Collectors.toSet());

            createAndAssert(new ExperimentItemsBatch(experimentItems), apiKey, workspaceName);

            // Make request with sorting parameter
            var sorting = SortingField.builder()
                    .field("data.expected_output")
                    .direction(Direction.ASC)
                    .build();
            var sortingParam = URLEncoder.encode(JsonUtils.writeValueAsString(List.of(sorting)),
                    StandardCharsets.UTF_8);
            var experimentIdsParam = JsonUtils.writeValueAsString(List.of(experiment.id()));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("page", 1)
                    .queryParam("size", 10)
                    .queryParam("experiment_ids", experimentIdsParam)
                    .queryParam("sorting", sortingParam)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualPage = actualResponse.readEntity(DatasetItemPage.class);
                assertThat(actualPage.content()).isNotEmpty();

                // With current test configuration (maxSizeToAllowSorting: -1),
                // dynamic sorting is always enabled, so sortableBy should be present
                assertThat(actualPage.sortableBy()).isNotNull();
                assertThat(actualPage.sortableBy()).isNotEmpty();

                /*
                 * If this test were run with maxSizeToAllowSorting set to 0:
                 *   assertThat(actualPage.sortableBy()).isEmpty();
                 *
                 * And the data would NOT be sorted by the requested field.
                 *
                 * The protection mechanism is:
                 *   1. Resource fetches ScopeMetadata
                 *   2. Checks canUseDynamicSorting()
                 *   3. If false: clears sorting fields and strips sortableBy from response
                 *   4. Frontend receives empty sortableBy and knows sorting is unavailable
                 */
            }
        }
    }

    @DisplayName("Get experiment items output columns test")
    @Nested
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class GetExperimentItemsOutputColumnsTest {

        private List<Trace> createTrace() {
            var trace = factory.manufacturePojo(Trace.class)
                    .toBuilder()
                    .metadata(JsonUtils
                            .getJsonNodeFromString(JsonUtils.writeValueAsString(Map.of("sql_cost", 10))))
                    .output(JsonUtils.readTree(Map.of(
                            "sql_tag", JsonUtils.readTree("sql_test"),
                            "sql_rate", JsonUtils.readTree(100),
                            "meta_field", JsonUtils.readTree(Map.of("version", new String[]{"10", "11", "12"})),
                            "releases", JsonUtils.readTree(
                                    List.of(
                                            Map.of("fixes", new String[]{"10", "11", "12"}, "version", "1.0"),
                                            Map.of("fixes", new String[]{"10", "11", "12"}, "version", "1.1"),
                                            Map.of("fixes", new String[]{"10", "45", "30"}, "version", "1.2"))),
                            "json_node", JsonUtils.readTree(Map.of("test", "1233", "test2", "12338")),
                            RandomStringUtils.randomAlphanumeric(5),
                            BigIntegerNode.valueOf(new BigInteger("18446744073709551615")),
                            RandomStringUtils.randomAlphanumeric(5), DoubleNode.valueOf(132432432.79995),
                            RandomStringUtils.randomAlphanumeric(5),
                            DoubleNode.valueOf(1.1844674407370955444555),
                            RandomStringUtils.randomAlphanumeric(5), IntNode.valueOf(100000000),
                            RandomStringUtils.randomAlphanumeric(5), BooleanNode.valueOf(true))))

                    .build();

            var trace2 = factory.manufacturePojo(Trace.class)
                    .toBuilder()
                    .metadata(JsonUtils
                            .getJsonNodeFromString(JsonUtils.writeValueAsString(Map.of("sql_cost", 10))))
                    .output(JsonUtils.readTree(Map.of(
                            "sql_tag", JsonUtils.readTree(Set.of(
                                    RandomStringUtils.randomAlphanumeric(5),
                                    RandomStringUtils.randomAlphanumeric(5),
                                    RandomStringUtils.randomAlphanumeric(5),
                                    RandomStringUtils.randomAlphanumeric(5),
                                    RandomStringUtils.randomAlphanumeric(5))),
                            "sql_rate", JsonUtils.readTree("100"),
                            "meta_field", JsonUtils.readTree("version"),
                            "json_node", JsonUtils.readTree(Map.of("test", "1233", "test2", "12338")),
                            "releases", DoubleNode.valueOf(132432432.79995))))
                    .build();

            var trace3 = factory.manufacturePojo(Trace.class)
                    .toBuilder()
                    .output(JsonUtils.readTree(RandomStringUtils.randomAlphanumeric(10)))
                    .build();

            var trace4 = factory.manufacturePojo(Trace.class)
                    .toBuilder()
                    .output(null)
                    .build();

            var trace5 = factory.manufacturePojo(Trace.class)
                    .toBuilder()
                    .output(JsonUtils.readTree(Set.of("test", "1233", "test2", "12338")))
                    .build();

            return List.of(trace, trace2, trace3, trace4, trace5);
        }

        @Test
        void getExperimentItemsOutputColumns__whenFetchingByDatasetId__thenReturnColumns() {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            String projectName = RandomStringUtils.randomAlphanumeric(20);

            List<Trace> traces = createTrace().stream()
                    .map(trace -> trace.toBuilder()
                            .projectName(projectName)
                            .build())
                    .toList();

            traces.parallelStream().forEach(trace -> createAndAssert(trace, workspaceName, apiKey));

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            // Creating 5 different experiment ids
            var experimentIds = IntStream.range(0, 5).mapToObj(__ -> GENERATOR.generate()).toList();

            List<ExperimentItem> experimentItems = IntStream.range(0, 5)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experimentIds.get(i))
                            .traceId(traces.get(i).id())
                            .datasetItemId(datasetItemBatch.items().get(i).id()).build())
                    .toList();

            var experimentItemsBatch = ExperimentItemsBatch.builder()
                    .experimentItems(Set.copyOf(experimentItems)).build();

            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            Set<Column> expectedOutput = getOutputDynamicColumns(traces);

            assertColumns(datasetId, apiKey, workspaceName, null, expectedOutput);

        }

        @Test
        void getExperimentItemsOutputColumns__whenFetchingByDatasetIdAndExperimentIds__thenReturnColumns() {
            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            String projectName = RandomStringUtils.randomAlphanumeric(20);

            List<Trace> traces = createTrace().stream()
                    .map(trace -> trace.toBuilder()
                            .projectName(projectName)
                            .build())
                    .toList();

            traces.parallelStream().forEach(trace -> createAndAssert(trace, workspaceName, apiKey));

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .build();

            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            // Creating 5 different experiment ids
            var experimentIds = IntStream.range(0, 2).mapToObj(__ -> GENERATOR.generate()).toList();

            List<ExperimentItem> experimentItems = IntStream.range(0, 2)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experimentIds.get(i))
                            .traceId(traces.get(i).id())
                            .datasetItemId(datasetItemBatch.items().get(i).id()).build())
                    .toList();

            var experimentItemsBatch = ExperimentItemsBatch.builder()
                    .experimentItems(Set.copyOf(experimentItems)).build();

            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            List<Trace> otherTraces = IntStream.range(0, 3)
                    .mapToObj(trace -> factory.manufacturePojo(Trace.class).toBuilder()
                            .projectName(projectName)
                            .build())
                    .toList();

            otherTraces.parallelStream().forEach(trace -> createAndAssert(trace, workspaceName, apiKey));

            var otherExperimentIds = IntStream.range(0, 3).mapToObj(__ -> GENERATOR.generate()).toList();

            List<DatasetItem> otherDatasetItems = datasetItemBatch.items().subList(2, 5);

            List<ExperimentItem> otherExperimentItems = IntStream.range(0, 3)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(otherExperimentIds.get(i))
                            .traceId(otherTraces.get(i).id())
                            .datasetItemId(otherDatasetItems.get(i).id()).build())
                    .toList();

            var otherExperimentItemsBatch = ExperimentItemsBatch.builder()
                    .experimentItems(Set.copyOf(otherExperimentItems)).build();

            createAndAssert(otherExperimentItemsBatch, apiKey, workspaceName);

            Set<Column> expectedOutput = getOutputDynamicColumns(traces);

            assertColumns(datasetId, apiKey, workspaceName, Set.copyOf(experimentIds), expectedOutput);
        }

        private void assertColumns(UUID datasetId, String apiKey, String workspaceName, Set<UUID> experimentIds,
                Set<Column> expectedOutput) {

            WebTarget request = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .path("output/columns");

            if (experimentIds != null) {
                request = request.queryParam("experiment_ids", JsonUtils.writeValueAsString(experimentIds));
            }

            try (var actualResponse = request
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_OK);
                assertThat(actualResponse.hasEntity()).isTrue();

                var actualColumns = actualResponse.readEntity(PageColumns.class);

                assertThat(actualColumns.columns())
                        .containsExactlyInAnyOrderElementsOf(expectedOutput);
            }
        }

    }

    private Set<Column> getOutputDynamicColumns(List<Trace> traces) {
        Map<String, List<JsonNode>> outputProperties = traces
                .stream()
                .map(Trace::output)
                .filter(Objects::nonNull)
                .filter(JsonNode::isObject)
                .map(JsonNode::fields)
                .flatMap(sourceIterator -> StreamSupport
                        .stream(Spliterators.spliteratorUnknownSize(sourceIterator, Spliterator.ORDERED), false))
                .collect(groupingBy(Map.Entry::getKey, mapping(Map.Entry::getValue, toList())));

        return outputProperties
                .entrySet()
                .stream()
                .map(entry -> {
                    Set<ColumnType> types = entry.getValue()
                            .stream()
                            .map(value -> getType(Map.entry(entry.getKey(), value)))
                            .collect(toSet());

                    return Column.builder().name(entry.getKey()).types(types).filterFieldPrefix("output").build();
                })
                .collect(toSet());
    }

    private Set<Column> getColumns(List<Map<String, JsonNode>> data) {

        HashSet<Column> columns = data
                .stream()
                .map(Map::entrySet)
                .flatMap(Collection::stream)
                .map(entry -> Column.builder()
                        .name(entry.getKey())
                        .types(Set.of(getType(entry)))
                        .filterFieldPrefix("data")
                        .build())
                .collect(toCollection(HashSet::new));

        Map<String, Set<ColumnType>> results = columns.stream()
                .collect(groupingBy(Column::name, mapping(Column::types, flatMapping(Set::stream, toSet()))));

        return results.entrySet()
                .stream()
                .map(entry -> Column.builder().name(entry.getKey()).types(entry.getValue()).filterFieldPrefix("data")
                        .build())
                .collect(toSet());
    }

    private ColumnType getType(Map.Entry<String, JsonNode> entry) {
        return switch (entry.getValue().getNodeType()) {
            case NUMBER -> ColumnType.NUMBER;
            case STRING -> ColumnType.STRING;
            case BOOLEAN -> ColumnType.BOOLEAN;
            case ARRAY -> ColumnType.ARRAY;
            case OBJECT -> ColumnType.OBJECT;
            case NULL -> ColumnType.NULL;
            default -> ColumnType.NULL;
        };
    }

    private void assertDatasetItemExperiments(DatasetItemPage actualPage, List<DatasetItem> datasetItems,
            List<ExperimentItem> experimentItems) {

        var actualDatasetItems = actualPage.content();
        assertPage(datasetItems, actualDatasetItems);
        assertThat(actualDatasetItems.size()).isEqualTo(experimentItems.size());

        for (int i = 0; i < actualDatasetItems.size(); i++) {
            var actualExperimentItems = actualDatasetItems.get(i).experimentItems();
            assertThat(actualExperimentItems).hasSize(1);
            assertThat(actualExperimentItems.getFirst())
                    .usingRecursiveComparison()
                    .ignoringFields(IGNORED_FIELDS_LIST)
                    .isEqualTo(experimentItems.get(i));

            var actualFeedbackScores = assertFeedbackScoresIgnoredFieldsAndSetThemToNull(
                    actualExperimentItems.getFirst(), USER).feedbackScores();

            if (ListUtils.emptyIfNull(experimentItems.get(i).feedbackScores()).isEmpty()) {
                assertThat(actualFeedbackScores).isNull();
                continue;
            }

            assertThat(actualFeedbackScores).hasSize(1);

            assertThat(actualFeedbackScores.getFirst())
                    .usingRecursiveComparison()
                    .withComparatorForType(BigDecimal::compareTo, BigDecimal.class)
                    .isEqualTo(experimentItems.get(i).feedbackScores().getFirst());
        }
    }

    private void putAndAssert(DatasetItemBatch batch, String workspaceName, String apiKey) {
        try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                .path("items")
                .request()
                .accept(MediaType.APPLICATION_JSON_TYPE)
                .header(HttpHeaders.AUTHORIZATION, apiKey)
                .header(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON)
                .header(WORKSPACE_HEADER, workspaceName)
                .put(Entity.json(batch))) {

            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(204);
            assertThat(actualResponse.hasEntity()).isFalse();
        }
    }

    private void createAndAssert(ExperimentItemsBatch request, String apiKey, String workspaceName) {
        experimentResourceClient.createExperimentItem(request.experimentItems(), apiKey, workspaceName);
    }

    private void createAndAssert(Trace trace, String workspaceName, String apiKey) {
        traceResourceClient.createTrace(trace, apiKey, workspaceName);
    }

    private String getTracesPath() {
        return URL_TEMPLATE_TRACES.formatted(baseURI);
    }

    private List<DatasetItem> getStreamedItems(Response response) {
        List<DatasetItem> items = new ArrayList<>();
        try (var inputStream = response.readEntity(new GenericType<ChunkedInput<String>>() {
        })) {
            while (true) {
                var item = inputStream.read();
                if (null == item) {
                    break;
                }
                items.add(JsonUtils.readValue(item, new TypeReference<DatasetItem>() {
                }));
            }
        }

        return items;
    }

    private void mockGetWorkspaceIdByName(String workspaceName, String workspaceId) {
        AuthTestUtils.mockGetWorkspaceIdByName(wireMock.server(), workspaceName, workspaceId);
    }

    private List<Dataset> prepareDatasetsListWithOnePublic() {
        var datasets = PodamFactoryUtils.manufacturePojoList(factory, Dataset.class).stream()
                .map(project -> project.toBuilder()
                        .visibility(PRIVATE)
                        .build())
                .collect(Collectors.toCollection(ArrayList::new));
        datasets.set(0, datasets.getFirst().toBuilder().visibility(PUBLIC).build());

        return datasets;
    }

    @Nested
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    @DisplayName("Get Dataset Experiment Items Stats")
    class GetDatasetExperimentItemsStats {

        @Test
        @DisplayName("Success: Get experiment items stats without filters")
        void getExperimentItemsStats__happyFlow() {
            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = datasetResourceClient.createDataset(dataset, apiKey, workspaceName);

            var experiment1 = experimentResourceClient.createPartialExperiment()
                    .datasetId(datasetId)
                    .build();
            createAndAssert(experiment1, apiKey, workspaceName);

            var experiment2 = experimentResourceClient.createPartialExperiment()
                    .datasetId(datasetId)
                    .build();
            createAndAssert(experiment2, apiKey, workspaceName);

            var datasetItem1 = factory.manufacturePojo(DatasetItem.class);
            var datasetItem2 = factory.manufacturePojo(DatasetItem.class);

            datasetResourceClient.createDatasetItems(
                    DatasetItemBatch.builder()
                            .items(List.of(datasetItem1, datasetItem2))
                            .datasetId(datasetId)
                            .build(),
                    workspaceName,
                    apiKey);

            var trace1 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment1.name())
                    .build();
            traceResourceClient.createTrace(trace1, apiKey, workspaceName);

            var experimentItem1 = ExperimentItem.builder()
                    .experimentId(experiment1.id())
                    .datasetItemId(datasetItem1.id())
                    .traceId(trace1.id())
                    .output(JsonUtils.readTree(Map.of("result", "output1")))
                    .build();

            var trace2 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment2.name())
                    .build();
            traceResourceClient.createTrace(trace2, apiKey, workspaceName);

            var experimentItem2 = ExperimentItem.builder()
                    .experimentId(experiment2.id())
                    .datasetItemId(datasetItem2.id())
                    .traceId(trace2.id())
                    .output(JsonUtils.readTree(Map.of("result", "output2")))
                    .build();

            var experimentItemsBatch = new ExperimentItemsBatch(Set.of(experimentItem1, experimentItem2));

            DatasetsResourceTest.this.createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Generate random feedback scores using PODAM
            var feedbackScore1 = factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder()
                    .id(trace1.id())
                    .name("accuracy")
                    .source(ScoreSource.SDK)
                    .build();

            var feedbackScore2 = factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder()
                    .id(trace2.id())
                    .name("accuracy")
                    .source(ScoreSource.SDK)
                    .build();

            var feedbackScores = List.of(feedbackScore1, feedbackScore2);
            traceResourceClient.feedbackScores(feedbackScores, apiKey, workspaceName);

            // Fetch traces to get actual durations
            var createdTrace1 = traceResourceClient.getById(trace1.id(), workspaceName, apiKey);
            var createdTrace2 = traceResourceClient.getById(trace2.id(), workspaceName, apiKey);
            var traceDurations = List.of(createdTrace1.duration(), createdTrace2.duration())
                    .stream()
                    .filter(Objects::nonNull)
                    .toList();

            // Calculate expected values from test data
            var experimentItemsCount = (long) experimentItemsBatch.experimentItems().size();
            var traceCount = (long) feedbackScores.size(); // One trace per feedback score
            var expectedAvgValue = feedbackScores.stream()
                    .map(FeedbackScoreItem::value)
                    .reduce(BigDecimal.ZERO, BigDecimal::add)
                    .divide(BigDecimal.valueOf(feedbackScores.size()), java.math.RoundingMode.HALF_UP);

            // Calculate duration percentiles from actual trace durations
            var durationPercentiles = StatsUtils.calculateQuantiles(traceDurations, List.of(0.50, 0.90, 0.99));

            var stats = datasetResourceClient.getDatasetExperimentItemsStats(
                    datasetId,
                    List.of(experiment1.id(), experiment2.id()),
                    apiKey,
                    workspaceName,
                    null);

            // Build complete expected stats list from calculated values
            List<ProjectStatItem<?>> expectedStats = List.of(
                    new CountValueStat(StatsMapper.EXPERIMENT_ITEMS_COUNT, experimentItemsCount),
                    new CountValueStat(StatsMapper.TRACE_COUNT, traceCount),
                    new AvgValueStat(StatsMapper.TOTAL_ESTIMATED_COST, 0.0), // No costs in test data
                    new PercentageValueStat(StatsMapper.DURATION,
                            new PercentageValues(
                                    durationPercentiles.get(0), // p50
                                    durationPercentiles.get(1), // p90
                                    durationPercentiles.get(2))), // p99
                    new AvgValueStat("feedback_scores.accuracy", expectedAvgValue.doubleValue()));

            // Assert the whole ProjectStats object using TraceAssertions
            TraceAssertions.assertStats(stats.stats(), expectedStats);
        }

        @Test
        @DisplayName("Success: Get experiment items stats with output filter")
        void getExperimentItemsStats__withOutputFilter() {
            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = datasetResourceClient.createDataset(dataset, apiKey, workspaceName);

            var experiment1 = experimentResourceClient.createPartialExperiment()
                    .datasetId(datasetId)
                    .build();
            createAndAssert(experiment1, apiKey, workspaceName);

            var datasetItem1 = factory.manufacturePojo(DatasetItem.class);
            var datasetItem2 = factory.manufacturePojo(DatasetItem.class);

            datasetResourceClient.createDatasetItems(
                    DatasetItemBatch.builder()
                            .items(List.of(datasetItem1, datasetItem2))
                            .datasetId(datasetId)
                            .build(),
                    workspaceName,
                    apiKey);

            var trace1 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment1.name())
                    .output(JsonUtils.readTree(Map.of("result", "success")))
                    .build();
            traceResourceClient.createTrace(trace1, apiKey, workspaceName);

            var experimentItem1 = ExperimentItem.builder()
                    .experimentId(experiment1.id())
                    .datasetItemId(datasetItem1.id())
                    .traceId(trace1.id())
                    .build();

            var trace2 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment1.name())
                    .output(JsonUtils.readTree(Map.of("result", "failure")))
                    .build();
            traceResourceClient.createTrace(trace2, apiKey, workspaceName);

            var experimentItem2 = ExperimentItem.builder()
                    .experimentId(experiment1.id())
                    .datasetItemId(datasetItem2.id())
                    .traceId(trace2.id())
                    .build();

            var experimentItemsBatch = new ExperimentItemsBatch(Set.of(experimentItem1, experimentItem2));

            DatasetsResourceTest.this.createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Generate random feedback scores using PODAM
            var feedbackScore1 = factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder()
                    .id(trace1.id())
                    .name("accuracy")
                    .source(ScoreSource.SDK)
                    .build();

            var feedbackScore2 = factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder()
                    .id(trace2.id())
                    .name("accuracy")
                    .source(ScoreSource.SDK)
                    .build();

            traceResourceClient.feedbackScores(List.of(feedbackScore1, feedbackScore2), apiKey,
                    workspaceName);

            // Fetch traces to get actual durations
            var createdTrace1 = traceResourceClient.getById(trace1.id(), workspaceName, apiKey);
            var createdTrace2 = traceResourceClient.getById(trace2.id(), workspaceName, apiKey);

            var outputFilter = List.of(ExperimentsComparisonFilter.builder()
                    .field(ExperimentsComparisonValidKnownField.OUTPUT.getQueryParamField())
                    .operator(Operator.CONTAINS)
                    .value("success")
                    .build());

            // Calculate expected values from test data
            // Expected: Only trace1 matches filter (output contains "success")
            var matchingExperimentItems = List.of(experimentItem1); // Only experimentItem1 has trace1 which has "success"
            var matchingFeedbackScores = List.of(feedbackScore1); // Only feedbackScore1 for trace1
            var experimentItemsCount = (long) matchingExperimentItems.size();
            var traceCount = (long) matchingFeedbackScores.size();
            var expectedAvgValue = matchingFeedbackScores.stream()
                    .map(FeedbackScoreItem::value)
                    .reduce(BigDecimal.ZERO, BigDecimal::add)
                    .divide(BigDecimal.valueOf(matchingFeedbackScores.size()), java.math.RoundingMode.HALF_UP);

            // Calculate duration percentiles from trace1 only (trace2 is filtered out)
            var traceDurations = List.of(createdTrace1.duration())
                    .stream()
                    .filter(Objects::nonNull)
                    .toList();
            var durationPercentiles = StatsUtils.calculateQuantiles(traceDurations, List.of(0.50, 0.90, 0.99));

            var stats = datasetResourceClient.getDatasetExperimentItemsStats(
                    datasetId,
                    List.of(experiment1.id()),
                    apiKey,
                    workspaceName,
                    outputFilter);

            // Assert the whole ProjectStats object - filter matches only trace1
            assertThat(stats).isNotNull();
            assertThat(stats.stats()).hasSizeGreaterThanOrEqualTo(4);

            // Verify required stat names are present
            assertThat(stats.stats())
                    .extracting(ProjectStatItem::getName)
                    .contains(StatsMapper.EXPERIMENT_ITEMS_COUNT, StatsMapper.TRACE_COUNT,
                            StatsMapper.TOTAL_ESTIMATED_COST, StatsMapper.DURATION, "feedback_scores.accuracy");

            // Build complete expected stats list from calculated values
            List<ProjectStatItem<?>> expectedStats = List.of(
                    new CountValueStat(StatsMapper.EXPERIMENT_ITEMS_COUNT, experimentItemsCount),
                    new CountValueStat(StatsMapper.TRACE_COUNT, traceCount),
                    new AvgValueStat(StatsMapper.TOTAL_ESTIMATED_COST, 0.0), // No costs in test data
                    new PercentageValueStat(StatsMapper.DURATION,
                            new PercentageValues(
                                    durationPercentiles.get(0), // p50
                                    durationPercentiles.get(1), // p90
                                    durationPercentiles.get(2))), // p99
                    new AvgValueStat("feedback_scores.accuracy", expectedAvgValue.doubleValue()));

            // Assert the whole ProjectStats object using TraceAssertions
            TraceAssertions.assertStats(stats.stats(), expectedStats);
        }

        @Test
        @DisplayName("Success: Get experiment items stats with multiple experiments")
        void getExperimentItemsStats__multipleExperiments() {
            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = datasetResourceClient.createDataset(dataset, apiKey, workspaceName);

            var experiment1 = experimentResourceClient.createPartialExperiment()
                    .datasetId(datasetId)
                    .build();
            createAndAssert(experiment1, apiKey, workspaceName);

            var experiment2 = experimentResourceClient.createPartialExperiment()
                    .datasetId(datasetId)
                    .build();
            createAndAssert(experiment2, apiKey, workspaceName);

            var experiment3 = experimentResourceClient.createPartialExperiment()
                    .datasetId(datasetId)
                    .build();
            createAndAssert(experiment3, apiKey, workspaceName);

            var datasetItem1 = factory.manufacturePojo(DatasetItem.class);
            var datasetItem2 = factory.manufacturePojo(DatasetItem.class);
            var datasetItem3 = factory.manufacturePojo(DatasetItem.class);

            datasetResourceClient.createDatasetItems(
                    DatasetItemBatch.builder()
                            .items(List.of(datasetItem1, datasetItem2, datasetItem3))
                            .datasetId(datasetId)
                            .build(),
                    workspaceName,
                    apiKey);

            var trace1 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment1.name())
                    .build();
            traceResourceClient.createTrace(trace1, apiKey, workspaceName);

            var trace2 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment2.name())
                    .build();
            traceResourceClient.createTrace(trace2, apiKey, workspaceName);

            var trace3 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment3.name())
                    .build();
            traceResourceClient.createTrace(trace3, apiKey, workspaceName);

            var experimentItem1 = ExperimentItem.builder()
                    .experimentId(experiment1.id())
                    .datasetItemId(datasetItem1.id())
                    .traceId(trace1.id())
                    .output(JsonUtils.readTree(Map.of("result", "output1")))
                    .build();

            var experimentItem2 = ExperimentItem.builder()
                    .experimentId(experiment2.id())
                    .datasetItemId(datasetItem2.id())
                    .traceId(trace2.id())
                    .output(JsonUtils.readTree(Map.of("result", "output2")))
                    .build();

            var experimentItem3 = ExperimentItem.builder()
                    .experimentId(experiment3.id())
                    .datasetItemId(datasetItem3.id())
                    .traceId(trace3.id())
                    .output(JsonUtils.readTree(Map.of("result", "output3")))
                    .build();

            var experimentItemsBatch = new ExperimentItemsBatch(
                    Set.of(experimentItem1, experimentItem2, experimentItem3));

            DatasetsResourceTest.this.createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Generate random feedback scores using PODAM
            var feedbackScore1 = factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder()
                    .id(trace1.id())
                    .name("quality")
                    .source(ScoreSource.SDK)
                    .build();

            var feedbackScore2 = factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder()
                    .id(trace2.id())
                    .name("quality")
                    .source(ScoreSource.SDK)
                    .build();

            var feedbackScore3 = factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder()
                    .id(trace3.id())
                    .name("quality")
                    .source(ScoreSource.SDK)
                    .build();

            var allExperimentItems = List.of(experimentItem1, experimentItem2, experimentItem3);
            var feedbackScores = List.of(feedbackScore1, feedbackScore2, feedbackScore3);

            traceResourceClient.feedbackScores(feedbackScores, apiKey, workspaceName);

            // Fetch traces to get actual durations
            var createdTrace1 = traceResourceClient.getById(trace1.id(), workspaceName, apiKey);
            var createdTrace2 = traceResourceClient.getById(trace2.id(), workspaceName, apiKey);
            var createdTrace3 = traceResourceClient.getById(trace3.id(), workspaceName, apiKey);

            // Calculate expected values from test data
            var experimentItemsCount = (long) allExperimentItems.size();
            var traceCount = (long) feedbackScores.size();
            var expectedAvgValue = feedbackScores.stream()
                    .map(FeedbackScoreItem::value)
                    .reduce(BigDecimal.ZERO, BigDecimal::add)
                    .divide(BigDecimal.valueOf(feedbackScores.size()), java.math.RoundingMode.HALF_UP);

            // Calculate duration percentiles from all three traces
            var traceDurations = List.of(createdTrace1.duration(), createdTrace2.duration(), createdTrace3.duration())
                    .stream()
                    .filter(Objects::nonNull)
                    .toList();
            var durationPercentiles = StatsUtils.calculateQuantiles(traceDurations, List.of(0.50, 0.90, 0.99));

            var stats = datasetResourceClient.getDatasetExperimentItemsStats(
                    datasetId,
                    List.of(experiment1.id(), experiment2.id(), experiment3.id()),
                    apiKey,
                    workspaceName,
                    null);

            // Build complete expected stats list from calculated values
            List<ProjectStatItem<?>> expectedStats = List.of(
                    new CountValueStat(StatsMapper.EXPERIMENT_ITEMS_COUNT, experimentItemsCount),
                    new CountValueStat(StatsMapper.TRACE_COUNT, traceCount),
                    new AvgValueStat(StatsMapper.TOTAL_ESTIMATED_COST, 0.0), // No costs in test data
                    new PercentageValueStat(StatsMapper.DURATION,
                            new PercentageValues(
                                    durationPercentiles.get(0), // p50
                                    durationPercentiles.get(1), // p90
                                    durationPercentiles.get(2))), // p99
                    new AvgValueStat("feedback_scores.quality", expectedAvgValue.doubleValue()));

            // Assert the whole ProjectStats object using TraceAssertions
            TraceAssertions.assertStats(stats.stats(), expectedStats);
        }

        @Test
        @DisplayName("Success: Get experiment items stats with is_not_empty feedback scores filter")
        void getExperimentItemsStats__withFeedbackScoresIsNotEmptyFilter() {
            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = datasetResourceClient.createDataset(dataset, apiKey, workspaceName);

            var experiment = experimentResourceClient.createPartialExperiment()
                    .datasetId(datasetId)
                    .build();
            createAndAssert(experiment, apiKey, workspaceName);

            // Create dataset items
            var datasetItem1 = factory.manufacturePojo(DatasetItem.class);
            var datasetItem2 = factory.manufacturePojo(DatasetItem.class);
            datasetResourceClient.createDatasetItems(
                    DatasetItemBatch.builder()
                            .items(List.of(datasetItem1, datasetItem2))
                            .datasetId(datasetId)
                            .build(),
                    workspaceName,
                    apiKey);

            // Create trace WITH feedback score
            var trace1 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment.name())
                    .build();
            traceResourceClient.createTrace(trace1, apiKey, workspaceName);

            // Create trace WITHOUT feedback score
            var trace2 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment.name())
                    .build();
            traceResourceClient.createTrace(trace2, apiKey, workspaceName);

            // Create experiment items
            var experimentItem1 = ExperimentItem.builder()
                    .experimentId(experiment.id())
                    .datasetItemId(datasetItem1.id())
                    .traceId(trace1.id())
                    .build();
            var experimentItem2 = ExperimentItem.builder()
                    .experimentId(experiment.id())
                    .datasetItemId(datasetItem2.id())
                    .traceId(trace2.id())
                    .build();

            var experimentItemsBatch = new ExperimentItemsBatch(Set.of(experimentItem1, experimentItem2));
            DatasetsResourceTest.this.createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Add feedback score to trace1 only
            var feedbackScoreName = "UserFeedback"; // Use name without spaces to avoid URL encoding issues
            var feedbackScore1 = factory.manufacturePojo(FeedbackScoreBatchItem.class).toBuilder()
                    .id(trace1.id())
                    .projectName(experiment.name())
                    .name(feedbackScoreName)
                    .value(BigDecimal.valueOf(4.5))
                    .build();
            traceResourceClient.feedbackScores(List.of(feedbackScore1), apiKey, workspaceName);

            // Create filter for is_not_empty on feedback score
            var filters = List.of(
                    ExperimentsComparisonFilter.builder()
                            .field(ExperimentsComparisonValidKnownField.FEEDBACK_SCORES.getQueryParamField())
                            .operator(Operator.IS_NOT_EMPTY)
                            .key(feedbackScoreName)
                            .value("")
                            .build());

            // Query with is_not_empty filter - should return only trace1 stats
            var stats = datasetResourceClient.getDatasetExperimentItemsStats(
                    datasetId,
                    List.of(experiment.id()),
                    apiKey,
                    workspaceName,
                    filters);

            // Fetch the trace with feedback score to get actual duration
            var createdTrace1 = traceResourceClient.getById(trace1.id(), workspaceName, apiKey);

            // Calculate expected values - should only include trace1 which has feedback score
            var experimentItemsCount = 1L; // Only experiment item 1 has feedback score
            var traceCount = 1L; // Only trace1 has feedback score
            var expectedAvgValue = feedbackScore1.value();

            var durationPercentiles = StatsUtils.calculateQuantiles(
                    List.of(createdTrace1.duration()),
                    List.of(0.50, 0.90, 0.99));

            // Build complete expected stats list
            List<ProjectStatItem<?>> expectedStats = List.of(
                    new CountValueStat(StatsMapper.EXPERIMENT_ITEMS_COUNT, experimentItemsCount),
                    new CountValueStat(StatsMapper.TRACE_COUNT, traceCount),
                    new AvgValueStat(StatsMapper.TOTAL_ESTIMATED_COST, 0.0), // No costs in test data
                    new PercentageValueStat(StatsMapper.DURATION,
                            new PercentageValues(
                                    durationPercentiles.get(0), // p50
                                    durationPercentiles.get(1), // p90
                                    durationPercentiles.get(2))), // p99
                    new AvgValueStat("feedback_scores.%s".formatted(feedbackScoreName),
                            expectedAvgValue.doubleValue()));

            // Assert the whole ProjectStats object using TraceAssertions
            TraceAssertions.assertStats(stats.stats(), expectedStats);
        }

        @ParameterizedTest(name = "Duration filter: {0} {1}ms")
        @MethodSource("durationFilterTestCases")
        @DisplayName("Success: Get experiment items with duration filter")
        void getExperimentItems__withDurationFilter__thenReturnFilteredItems(
                Operator operator,
                long filterValueMs,
                List<Integer> expectedTraceIndices) {

            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = datasetResourceClient.createDataset(dataset, apiKey, workspaceName);

            var experiment = experimentResourceClient.createPartialExperiment()
                    .datasetId(datasetId)
                    .build();
            createAndAssert(experiment, apiKey, workspaceName);

            // Create dataset items
            var datasetItem1 = factory.manufacturePojo(DatasetItem.class);
            var datasetItem2 = factory.manufacturePojo(DatasetItem.class);
            var datasetItem3 = factory.manufacturePojo(DatasetItem.class);
            datasetResourceClient.createDatasetItems(
                    DatasetItemBatch.builder()
                            .items(List.of(datasetItem1, datasetItem2, datasetItem3))
                            .datasetId(datasetId)
                            .build(),
                    workspaceName,
                    apiKey);

            // Create traces with KNOWN, CONTROLLED durations
            // Using explicit start and end times to ensure predictable durations
            var baseTime = Instant.now().truncatedTo(ChronoUnit.SECONDS);

            // Trace 1: Duration = 500ms
            var trace1 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment.name())
                    .startTime(baseTime)
                    .endTime(baseTime.plusMillis(500))
                    .build();
            traceResourceClient.createTrace(trace1, apiKey, workspaceName);

            // Trace 2: Duration = 1500ms
            var trace2 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment.name())
                    .startTime(baseTime)
                    .endTime(baseTime.plusMillis(1500))
                    .build();
            traceResourceClient.createTrace(trace2, apiKey, workspaceName);

            // Trace 3: Duration = 2500ms
            var trace3 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(experiment.name())
                    .startTime(baseTime)
                    .endTime(baseTime.plusMillis(2500))
                    .build();
            traceResourceClient.createTrace(trace3, apiKey, workspaceName);

            // Create experiment items linking dataset items to traces
            var experimentItem1 = ExperimentItem.builder()
                    .experimentId(experiment.id())
                    .datasetItemId(datasetItem1.id())
                    .traceId(trace1.id())
                    .build();
            var experimentItem2 = ExperimentItem.builder()
                    .experimentId(experiment.id())
                    .datasetItemId(datasetItem2.id())
                    .traceId(trace2.id())
                    .build();
            var experimentItem3 = ExperimentItem.builder()
                    .experimentId(experiment.id())
                    .datasetItemId(datasetItem3.id())
                    .traceId(trace3.id())
                    .build();

            var experimentItemsBatch = new ExperimentItemsBatch(
                    Set.of(experimentItem1, experimentItem2, experimentItem3));
            DatasetsResourceTest.this.createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Apply duration filter
            var filters = List.of(
                    ExperimentsComparisonFilter.builder()
                            .field(ExperimentsComparisonValidKnownField.DURATION.getQueryParamField())
                            .operator(operator)
                            .value(String.valueOf(filterValueMs))
                            .build());

            // Query stats with duration filter
            var stats = datasetResourceClient.getDatasetExperimentItemsStats(
                    datasetId,
                    List.of(experiment.id()),
                    apiKey,
                    workspaceName,
                    filters);

            // Verify experiment items count matches expected filtered results
            var expectedItemsCount = (long) expectedTraceIndices.size();

            // Assert: Verify experiment items count and trace count match expected filtered results
            // Using direct assertion since we only care about the count matching the expected filter results
            assertThat(stats.stats()).isNotNull();

            var experimentItemsCountStat = stats.stats().stream()
                    .filter(stat -> StatsMapper.EXPERIMENT_ITEMS_COUNT.equals(stat.getName()))
                    .findFirst()
                    .orElseThrow(() -> new AssertionError("Expected EXPERIMENT_ITEMS_COUNT stat to be present"));

            assertThat(((CountValueStat) experimentItemsCountStat).getValue())
                    .as("Duration filter %s %dms should return %d items (traces %s)",
                            operator, filterValueMs, expectedItemsCount, expectedTraceIndices)
                    .isEqualTo(expectedItemsCount);

            // Verify trace count matches expected filtered results
            var traceCountStat = stats.stats().stream()
                    .filter(stat -> StatsMapper.TRACE_COUNT.equals(stat.getName()))
                    .findFirst()
                    .orElseThrow(() -> new AssertionError("Expected TRACE_COUNT stat to be present"));

            assertThat(((CountValueStat) traceCountStat).getValue())
                    .isEqualTo(expectedItemsCount);
        }

        /**
         * Test cases for duration filtering with different operators.
         * Parameters: operator, filter value in ms, expected trace indices (0-based)
         *
         * Test data:
         * - Trace 0: 500ms
         * - Trace 1: 1500ms
         * - Trace 2: 2500ms
         */
        static Stream<Arguments> durationFilterTestCases() {
            return Stream.of(
                    // GREATER_THAN: Return traces with duration > threshold
                    Arguments.of(Operator.GREATER_THAN, 1000L, List.of(1, 2)), // > 1000ms: 1500ms, 2500ms
                    Arguments.of(Operator.GREATER_THAN, 2000L, List.of(2)), // > 2000ms: 2500ms
                    Arguments.of(Operator.GREATER_THAN, 3000L, List.of()), // > 3000ms: none

                    // GREATER_THAN_EQUAL: Return traces with duration >= threshold
                    Arguments.of(Operator.GREATER_THAN_EQUAL, 1500L, List.of(1, 2)), // >= 1500ms: 1500ms, 2500ms
                    Arguments.of(Operator.GREATER_THAN_EQUAL, 2500L, List.of(2)), // >= 2500ms: 2500ms

                    // LESS_THAN: Return traces with duration < threshold
                    Arguments.of(Operator.LESS_THAN, 1000L, List.of(0)), // < 1000ms: 500ms
                    Arguments.of(Operator.LESS_THAN, 2000L, List.of(0, 1)), // < 2000ms: 500ms, 1500ms
                    Arguments.of(Operator.LESS_THAN, 400L, List.of()), // < 400ms: none

                    // LESS_THAN_EQUAL: Return traces with duration <= threshold
                    Arguments.of(Operator.LESS_THAN_EQUAL, 500L, List.of(0)), // <= 500ms: 500ms
                    Arguments.of(Operator.LESS_THAN_EQUAL, 1500L, List.of(0, 1)), // <= 1500ms: 500ms, 1500ms

                    // EQUAL: Return traces with duration = threshold
                    Arguments.of(Operator.EQUAL, 500L, List.of(0)), // = 500ms: 500ms
                    Arguments.of(Operator.EQUAL, 1500L, List.of(1)), // = 1500ms: 1500ms
                    Arguments.of(Operator.EQUAL, 1000L, List.of()), // = 1000ms: none

                    // NOT_EQUAL: Return traces with duration != threshold
                    Arguments.of(Operator.NOT_EQUAL, 1500L, List.of(0, 2)), // != 1500ms: 500ms, 2500ms
                    Arguments.of(Operator.NOT_EQUAL, 1000L, List.of(0, 1, 2)) // != 1000ms: all
            );
        }
    }

    @Nested
    @DisplayName("OPIK-2469: Cross-Project Traces Duplicate Test")
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    class CrossProjectTracesDuplicateTest {

        @Test
        @DisplayName("Should return unique experiment items when trace has spans in multiple projects")
        void findDatasetItemsWithExperimentItems__whenTraceHasSpansInMultipleProjects__thenReturnUniqueItems() {

            var workspaceName = UUID.randomUUID().toString();
            var apiKey = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            // Create dataset
            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            // Create dataset item
            var datasetItem = factory.manufacturePojo(DatasetItem.class);
            var datasetItemBatch = DatasetItemBatch.builder()
                    .datasetId(datasetId)
                    .items(List.of(datasetItem))
                    .build();
            putAndAssert(datasetItemBatch, workspaceName, apiKey);

            // Create Project A and Project B
            var projectA = UUID.randomUUID().toString();
            var projectB = UUID.randomUUID().toString();

            // Create trace in Project A
            var trace1 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(projectA)
                    .build();
            createAndAssert(trace1, workspaceName, apiKey);

            // Create span in Project A for trace1
            var span1InProjectA = factory.manufacturePojo(Span.class).toBuilder()
                    .projectName(projectA)
                    .traceId(trace1.id())
                    .build();
            createSpan(span1InProjectA, apiKey, workspaceName);

            // ROOT CAUSE SIMULATION: Create span in Project B for the SAME trace (trace1)
            // This creates a cross-project trace scenario through the API
            var span1InProjectB = factory.manufacturePojo(Span.class).toBuilder()
                    .projectName(projectB)
                    .traceId(trace1.id())
                    .build();
            createSpan(span1InProjectB, apiKey, workspaceName);

            // Create another trace in Project A (no cross-project issue)
            var trace2 = factory.manufacturePojo(Trace.class).toBuilder()
                    .projectName(projectA)
                    .build();
            createAndAssert(trace2, workspaceName, apiKey);

            // Create experiment items for both traces
            var experimentId = GENERATOR.generate();
            var experimentItem1 = factory.manufacturePojo(ExperimentItem.class).toBuilder()
                    .experimentId(experimentId)
                    .datasetItemId(datasetItem.id())
                    .traceId(trace1.id())
                    .input(trace1.input())
                    .output(trace1.output())
                    .build();

            var experimentItem2 = factory.manufacturePojo(ExperimentItem.class).toBuilder()
                    .experimentId(experimentId)
                    .datasetItemId(datasetItem.id())
                    .traceId(trace2.id())
                    .input(trace2.input())
                    .output(trace2.output())
                    .build();

            var experimentItemsBatch = ExperimentItemsBatch.builder()
                    .experimentItems(Set.of(experimentItem1, experimentItem2))
                    .build();
            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Query the endpoint
            var result = datasetResourceClient.getDatasetItemsWithExperimentItems(
                    datasetId,
                    List.of(experimentId),
                    apiKey,
                    workspaceName);

            // Assert results
            assertThat(result).isNotNull();
            assertThat(result.content()).hasSize(1);

            var datasetItemResult = result.content().get(0);
            assertThat(datasetItemResult.id()).isEqualTo(datasetItem.id());

            // CRITICAL ASSERTION: Should have exactly 2 unique experiment items (no duplicates)
            // Without the fix, trace1 appears twice because it has spans in 2 projects
            var experimentItems = datasetItemResult.experimentItems();
            assertThat(experimentItems).isNotNull();

            // Count experiment items by their ID to detect duplicates
            var experimentItemIds = experimentItems.stream()
                    .map(ExperimentItem::id)
                    .collect(Collectors.toList());

            var uniqueIds = new HashSet<>(experimentItemIds);

            // THIS IS THE KEY ASSERTION - Verifies fix for OPIK-2469
            assertThat(experimentItemIds)
                    .as("Should not contain duplicate experiment item IDs - trace1 has spans in 2 projects but should appear once")
                    .hasSameSizeAs(uniqueIds)
                    .as("Should have exactly 2 unique experiment items")
                    .hasSize(2);

            // Verify the correct experiment items are present
            assertThat(uniqueIds).containsExactlyInAnyOrder(experimentItem1.id(), experimentItem2.id());

            // Verify each experiment item appears only once
            experimentItemIds.forEach(id -> {
                long count = experimentItemIds.stream().filter(i -> i.equals(id)).count();
                assertThat(count)
                        .as("Experiment item '%s' should appear exactly once, but appears '%d' times", id, count)
                        .isEqualTo(1);
            });
        }
    }

    @Nested
    @TestInstance(TestInstance.Lifecycle.PER_CLASS)
    @DisplayName("Experiment Items Sorting and Filtering: OPIK-2936")
    class ExperimentItemsSortingAndFiltering {

        @Test
        @DisplayName("should sort experiment items by total_estimated_cost in descending order")
        void sortByTotalEstimatedCost__whenDescendingOrder__thenReturnSorted() {
            var apiKey = UUID.randomUUID().toString();
            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            // Create project name for traces
            var projectName = RandomStringUtils.randomAlphanumeric(10);

            // Create dataset
            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            // Create 3 dataset items
            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .items(IntStream.range(0, 3)
                            .mapToObj(i -> factory.manufacturePojo(DatasetItem.class))
                            .toList())
                    .build();
            putAndAssert(datasetItemBatch, workspaceName, apiKey);
            var datasetItems = datasetItemBatch.items();

            // Create traces and spans with costs
            var costValues = List.of(
                    BigDecimal.valueOf(10.50),
                    BigDecimal.valueOf(25.75),
                    BigDecimal.valueOf(5.25));

            // Create traces using PodamFactoryUtils
            var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream()
                    .limit(3)
                    .map(trace -> trace.toBuilder().projectName(projectName).build())
                    .toList();

            var traceIds = traces.stream()
                    .map(trace -> createTrace(trace, apiKey, workspaceName))
                    .toList();

            // Create spans in batch
            var spans = IntStream.range(0, 3)
                    .mapToObj(i -> factory.manufacturePojo(Span.class).toBuilder()
                            .projectName(projectName)
                            .traceId(traceIds.get(i))
                            .totalEstimatedCost(costValues.get(i))
                            .build())
                    .toList();

            spanResourceClient.batchCreateSpans(spans, apiKey, workspaceName);

            // Create experiment
            var experimentId = GENERATOR.generate();

            // Create experiment items linked to traces
            var experimentItems = IntStream.range(0, 3)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experimentId)
                            .datasetItemId(datasetItems.get(i).id())
                            .traceId(traceIds.get(i))
                            .build())
                    .toList();

            var experimentItemsBatch = ExperimentItemsBatch.builder()
                    .experimentItems(new HashSet<>(experimentItems))
                    .build();
            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Query with sorting by total_estimated_cost DESC
            var sorting = toURLEncodedQueryParam(
                    List.of(new SortingField("total_estimated_cost", Direction.DESC)));
            var experimentIdsParam = JsonUtils.writeValueAsString(List.of(experimentId));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("experiment_ids", experimentIdsParam)
                    .queryParam("sorting", sorting)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_OK);
                var actualPage = actualResponse.readEntity(DatasetItemPage.class);

                assertThat(actualPage.content()).hasSize(3);

                // Verify sorting order: 25.75, 10.50, 5.25
                var costs = actualPage.content().stream()
                        .map(item -> item.experimentItems().get(0).totalEstimatedCost())
                        .toList();

                assertThat(costs).hasSize(3);
                assertThat(costs.get(0)).isEqualByComparingTo(BigDecimal.valueOf(25.75));
                assertThat(costs.get(1)).isEqualByComparingTo(BigDecimal.valueOf(10.50));
                assertThat(costs.get(2)).isEqualByComparingTo(BigDecimal.valueOf(5.25));
            }
        }

        @Test
        @DisplayName("should sort experiment items by usage.total_tokens in ascending order")
        void sortByUsageTotalTokens__whenAscendingOrder__thenReturnSorted() {
            var apiKey = UUID.randomUUID().toString();
            var workspaceName = UUID.randomUUID().toString();
            var workspaceId = UUID.randomUUID().toString();

            mockTargetWorkspace(apiKey, workspaceName, workspaceId);

            // Create project name for traces
            var projectName = RandomStringUtils.randomAlphanumeric(10);

            // Create dataset
            var dataset = factory.manufacturePojo(Dataset.class);
            var datasetId = createAndAssert(dataset, apiKey, workspaceName);

            // Create 3 dataset items
            var datasetItemBatch = factory.manufacturePojo(DatasetItemBatch.class).toBuilder()
                    .datasetId(datasetId)
                    .items(IntStream.range(0, 3)
                            .mapToObj(i -> factory.manufacturePojo(DatasetItem.class))
                            .toList())
                    .build();
            putAndAssert(datasetItemBatch, workspaceName, apiKey);
            var datasetItems = datasetItemBatch.items();

            // Create traces and spans with usage data
            var usageValues = List.of(
                    Map.of("total_tokens", 150, "prompt_tokens", 100, "completion_tokens", 50),
                    Map.of("total_tokens", 50, "prompt_tokens", 30, "completion_tokens", 20),
                    Map.of("total_tokens", 100, "prompt_tokens", 60, "completion_tokens", 40));

            // Create traces using PodamFactoryUtils
            var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class).stream()
                    .limit(3)
                    .map(trace -> trace.toBuilder().projectName(projectName).build())
                    .toList();

            var traceIds = traces.stream()
                    .map(trace -> createTrace(trace, apiKey, workspaceName))
                    .toList();

            // Create spans in batch
            var spans = IntStream.range(0, 3)
                    .mapToObj(i -> factory.manufacturePojo(Span.class).toBuilder()
                            .projectName(projectName)
                            .traceId(traceIds.get(i))
                            .usage(usageValues.get(i))
                            .build())
                    .toList();

            spanResourceClient.batchCreateSpans(spans, apiKey, workspaceName);

            // Create experiment
            var experimentId = GENERATOR.generate();

            // Create experiment items linked to traces
            var experimentItems = IntStream.range(0, 3)
                    .mapToObj(i -> factory.manufacturePojo(ExperimentItem.class).toBuilder()
                            .experimentId(experimentId)
                            .datasetItemId(datasetItems.get(i).id())
                            .traceId(traceIds.get(i))
                            .build())
                    .toList();

            var experimentItemsBatch = ExperimentItemsBatch.builder()
                    .experimentItems(new HashSet<>(experimentItems))
                    .build();
            createAndAssert(experimentItemsBatch, apiKey, workspaceName);

            // Query with sorting by usage.total_tokens ASC
            var sorting = toURLEncodedQueryParam(
                    List.of(new SortingField("usage.total_tokens", Direction.ASC)));
            var experimentIdsParam = JsonUtils.writeValueAsString(List.of(experimentId));

            try (var actualResponse = client.target(BASE_RESOURCE_URI.formatted(baseURI))
                    .path(datasetId.toString())
                    .path(DATASET_ITEMS_WITH_EXPERIMENT_ITEMS_PATH)
                    .queryParam("experiment_ids", experimentIdsParam)
                    .queryParam("sorting", sorting)
                    .request()
                    .header(HttpHeaders.AUTHORIZATION, apiKey)
                    .header(WORKSPACE_HEADER, workspaceName)
                    .get()) {

                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_OK);
                var actualPage = actualResponse.readEntity(DatasetItemPage.class);

                assertThat(actualPage.content()).hasSize(3);

                // Verify sorting order: 50, 100, 150
                var tokens = actualPage.content().stream()
                        .map(item -> item.experimentItems().get(0).usage().get("total_tokens"))
                        .toList();

                assertThat(tokens).containsExactly(50L, 100L, 150L);
            }
        }
    }
}
